Propagate Future cancellations to getStageAsyncExecution supplied stages

Fixes #290

Author: jhalterman

src/main/java/net/jodah/failsafe/FailsafeExecutor.java

@@ -146,6 +146,8 @@ public <T extends R> CompletableFuture<T> getAsyncExecution(AsyncRunnable<T> run
   /**
    * Executes the {@code supplier} asynchronously until the resulting future is successfully completed or the configured
    * policies are exceeded.
+   * <p>Cancelling the resulting {@link CompletableFuture} will automatically cancels the supplied {@link
+   * CompletionStage} if it's a {@link Future}.</p>
    * <p>
    * If the execution fails because a {@link Timeout} is exceeded, the resulting future is completed exceptionally with
    * {@link TimeoutExceededException}.
@@ -165,6 +167,8 @@ public <T extends R> CompletableFuture<T> getStageAsync(CheckedSupplier<? extend
   /**
    * Executes the {@code supplier} asynchronously until the resulting future is successfully completed or the configured
    * policies are exceeded.
+   * <p>Cancelling the resulting {@link CompletableFuture} will automatically cancels the supplied {@link
+   * CompletionStage} if it's a {@link Future}.</p>
    * <p>
    * If the execution fails because a {@link Timeout} is exceeded, the resulting future is completed exceptionally with
    * {@link TimeoutExceededException}.
@@ -191,6 +195,8 @@ public <T extends R> CompletableFuture<T> getStageAsync(
    * completed. Any exception that is thrown from the {@code supplier} will automatically be recorded via {@code
    * AsyncExecution.recordFailure}.
    * </p>
+   * <p>Cancelling the resulting {@link CompletableFuture} will automatically cancels the supplied {@link
+   * CompletionStage} if it's a {@link Future}.</p>
    * <p>
    * If the execution fails because a {@link Timeout} is exceeded, the resulting future is completed exceptionally with
    * {@link TimeoutExceededException}.

src/main/java/net/jodah/failsafe/FailsafeFuture.java

@@ -21,6 +21,7 @@
 import java.util.Map.Entry;
 import java.util.concurrent.CancellationException;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.Future;
 import java.util.function.BiConsumer;
 
 /**
@@ -39,7 +40,7 @@ public class FailsafeFuture<R> extends CompletableFuture<R> {
   private AbstractExecution<R> newestExecution;
   Map<Integer, BiConsumer<Boolean, ExecutionResult>> cancelFunctions;
   // Whether a cancel with interrupt has already occurred
-  boolean cancelledWithInterrupt;
+  private boolean cancelledWithInterrupt;
 
   FailsafeFuture(FailsafeExecutor<R> executor) {
     this.executor = executor;
@@ -135,4 +136,15 @@ synchronized void injectCancelFn(int policyIndex, BiConsumer<Boolean, ExecutionR
       cancelFunctions = new HashMap<>();
     cancelFunctions.put(policyIndex, cancelFn);
   }
+
+  /**
+   * Propogates any previous cancellation to the {@code future}, either by cancelling it immediately or storing a cancel
+   * function for later.
+   */
+  synchronized void propagateCancellation(Future<R> future, int policyIndex) {
+    if (isCancelled())
+      future.cancel(cancelledWithInterrupt);
+    else
+      injectCancelFn(policyIndex, (mayInterrupt, cancelResult) -> future.cancel(mayInterrupt));
+  }
 }

src/main/java/net/jodah/failsafe/Functions.java

@@ -132,15 +132,8 @@ static <R> Function<AsyncExecution<R>, CompletableFuture<ExecutionResult>> getPr
         CompletionStage<? extends R> stage = supplier.get(execution);
 
         // Propagate outer cancellations to the stage
-        if (stage instanceof Future) {
-          Future<R> future = (Future<R>) stage;
-          synchronized (execution.future) {
-            if (execution.future.isCancelled())
-              future.cancel(execution.future.cancelledWithInterrupt);
-            else
-              execution.future.injectCancelFn(-1, (mayInterrupt, cancelResult) -> future.cancel(mayInterrupt));
-          }
-        }
+        if (stage instanceof Future)
+          execution.future.propagateCancellation((Future<R>) stage, -1);
 
         stage.whenComplete((result, failure) -> {
           if (failure instanceof CompletionException)
@@ -165,6 +158,7 @@ static <R> Function<AsyncExecution<R>, CompletableFuture<ExecutionResult>> getPr
    *
    * @param <R> result type
    */
+  @SuppressWarnings("unchecked")
   static <R> Function<AsyncExecution<R>, CompletableFuture<ExecutionResult>> getPromiseOfStageExecution(
     AsyncSupplier<R, ? extends CompletionStage<? extends R>> supplier) {
 
@@ -175,6 +169,11 @@ static <R> Function<AsyncExecution<R>, CompletableFuture<ExecutionResult>> getPr
         execution.preExecute();
         asyncFutureLock.acquire();
         CompletionStage<? extends R> stage = supplier.get(execution);
+
+        // Propagate outer cancellations to the stage
+        if (stage instanceof Future)
+          execution.future.propagateCancellation((Future<R>) stage, -1);
+
         stage.whenComplete((innerResult, failure) -> {
           try {
             if (failure != null)

src/test/java/net/jodah/failsafe/functional/FailsafeFutureCancellationTest.java

@@ -3,20 +3,22 @@
 import net.jodah.concurrentunit.Waiter;
 import net.jodah.failsafe.*;
 import net.jodah.failsafe.event.ExecutionCompletedEvent;
+import net.jodah.failsafe.function.AsyncSupplier;
 import net.jodah.failsafe.function.CheckedSupplier;
 import org.testng.annotations.Test;
 
 import java.time.Duration;
 import java.util.concurrent.*;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.BiConsumer;
 
 import static org.testng.Assert.*;
 
 /**
- * Tests behavior when a Failsafe Future is manually cancelled.
+ * Tests behavior when a FailsafeFuture is manually cancelled.
  */
 @Test
-public class ManualCancellationTest extends Testing {
+public class FailsafeFutureCancellationTest extends Testing {
   /**
    * Asserts that cancelling a FailsafeFuture causes both retry policies to stop.
    */
@@ -55,29 +57,62 @@ public void testCancelWithNestedRetries() throws Throwable {
   /**
    * Asserts that FailsafeFuture cancellations are propagated to a CompletionStage.
    */
-  public void shouldPropagateCancellationToCompletionStage() throws Throwable {
+  public void shouldPropagateCancellationToStage() throws Throwable {
     // Given
     Policy<String> retryPolicy = new RetryPolicy<>();
-    Waiter cancelledWaiter = new Waiter();
-    CheckedSupplier<CompletionStage<String>> computeSomething = () -> {
-      CompletableFuture<String> future = new CompletableFuture<>();
-      future.whenComplete((r, t) -> {
-        if (t instanceof CancellationException)
-          cancelledWaiter.resume();
-      });
-      return future;
+    Waiter cancellationWaiter = new Waiter();
+    BiConsumer<String, Throwable> waiterResumer = (r, t) -> {
+      if (t instanceof CancellationException)
+        cancellationWaiter.resume();
+    };
+    CheckedSupplier<CompletionStage<String>> doWork = () -> {
+      CompletableFuture<String> promise = new CompletableFuture<>();
+      promise.whenComplete(waiterResumer);
+
+      // Simulate asynchronous work
+      runInThread(() -> Thread.sleep(1000));
+      return promise;
     };
 
     // When
-    CompletableFuture<String> future = Failsafe.with(retryPolicy).getStageAsync(computeSomething);
-    future.whenComplete((r, t) -> {
+    CompletableFuture<String> failsafeFuture = Failsafe.with(retryPolicy).getStageAsync(doWork);
+    failsafeFuture.whenComplete(waiterResumer);
+    failsafeFuture.cancel(true);
+
+    // Then
+    Asserts.assertThrows(failsafeFuture::get, CancellationException.class);
+    // Wait for the promise and failsafeFuture to complete with cancellation
+    cancellationWaiter.await(1000, 2);
+  }
+
+  /**
+   * Asserts that FailsafeFuture cancellations are propagated to a CompletionStage in an async integration execution.
+   */
+  public void shouldPropagateCancellationToStageAsyncExecution() throws Throwable {
+    // Given
+    Policy<String> retryPolicy = new RetryPolicy<>();
+    Waiter cancellationWaiter = new Waiter();
+    BiConsumer<String, Throwable> waiterResumer = (r, t) -> {
       if (t instanceof CancellationException)
-        cancelledWaiter.resume();
-    });
-    future.cancel(true);
+        cancellationWaiter.resume();
+    };
+    AsyncSupplier<String, CompletionStage<String>> doWork = exec -> {
+      CompletableFuture<String> promise = new CompletableFuture<>();
+      promise.whenComplete(waiterResumer);
+
+      // Simulate asynchronous work
+      runInThread(() -> Thread.sleep(1000));
+      return promise;
+    };
+
+    // When
+    CompletableFuture<String> failsafeFuture = Failsafe.with(retryPolicy).getStageAsyncExecution(doWork);
+    failsafeFuture.whenComplete(waiterResumer);
+    failsafeFuture.cancel(true);
 
     // Then
-    Asserts.assertThrows(future::get, CancellationException.class);
-    cancelledWaiter.await(1000, 2);
+    Asserts.assertThrows(failsafeFuture::get, CancellationException.class);
+    // Wait for the promise and failsafeFuture to complete with cancellation
+    cancellationWaiter.await(1000, 2);
   }
 }