Cancel inner futures returned by getStageAsync

Fixes #266

Author: jhalterman

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

@@ -182,7 +182,7 @@ void executeAsync(boolean asyncExecution) {
     if (!asyncExecution)
       outerExecutionSupplier.get().whenComplete(this::complete);
     else
-      future.inject(scheduler.schedule(innerExecutionSupplier::get, 0, TimeUnit.NANOSECONDS));
+      future.injectPolicy(scheduler.schedule(innerExecutionSupplier::get, 0, TimeUnit.NANOSECONDS));
   }
 
   /**

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

@@ -34,9 +34,10 @@ public class FailsafeFuture<T> extends CompletableFuture<T> {
   private AbstractExecution execution;
 
   // Mutable state, guarded by "this"
-  private Future<T> dependency;
+  private Future<T> policyExecFuture;
+  private Future<?> dependentStageFuture;
   private Runnable cancelFn;
-  private List<Future<T>> timeoutDependencies;
+  private List<Future<T>> timeoutFutures;
 
   FailsafeFuture(FailsafeExecutor<T> executor) {
     this.executor = executor;
@@ -96,11 +97,11 @@ synchronized boolean completeResult(ExecutionResult result) {
   }
 
   synchronized Future<T> getDependency() {
-    return dependency;
+    return policyExecFuture;
   }
 
   synchronized List<Future<T>> getTimeoutDelegates() {
-    return timeoutDependencies;
+    return timeoutFutures;
   }
 
   /**
@@ -109,14 +110,16 @@ synchronized List<Future<T>> getTimeoutDelegates() {
    */
   synchronized boolean cancelDependencies(boolean interruptDelegate, boolean result) {
     execution.interrupted = interruptDelegate;
-    if (dependency != null)
-      result = dependency.cancel(interruptDelegate);
+    if (policyExecFuture != null)
+      result = policyExecFuture.cancel(interruptDelegate);
+    if (dependentStageFuture != null)
+      dependentStageFuture.cancel(interruptDelegate);
     if (cancelFn != null)
       cancelFn.run();
-    if (timeoutDependencies != null) {
-      for (Future<T> timeoutDelegate : timeoutDependencies)
+    if (timeoutFutures != null) {
+      for (Future<T> timeoutDelegate : timeoutFutures)
         timeoutDelegate.cancel(false);
-      timeoutDependencies.clear();
+      timeoutFutures.clear();
     }
     return result;
   }
@@ -126,11 +129,18 @@ void inject(AbstractExecution execution) {
   }
 
   /**
-   * Injects a {@code dependency} to be cancelled when this future is cancelled.
+   * Injects a {@code policyExecFuture} to be cancelled when this future is cancelled.
    */
   @SuppressWarnings({ "unchecked", "rawtypes" })
-  synchronized void inject(Future<?> dependency) {
-    this.dependency = (Future) dependency;
+  synchronized void injectPolicy(Future<?> policyExecFuture) {
+    this.policyExecFuture = (Future) policyExecFuture;
+  }
+
+  /**
+   * Injects a {@code dependentStageFuture} to be cancelled when this future is cancelled.
+   */
+  synchronized void injectStage(Future<?> dependentStageFuture) {
+    this.dependentStageFuture = dependentStageFuture;
   }
 
   /**
@@ -141,12 +151,12 @@ synchronized void injectCancelFn(Runnable cancelFn) {
   }
 
   /**
-   * Injects a {@code timeoutDependency} to be cancelled when this future is cancelled.
+   * Injects a {@code scheduledTimeoutExec} to be cancelled when this future is cancelled.
    */
   @SuppressWarnings({ "unchecked", "rawtypes" })
-  synchronized void injectTimeout(Future<?> timeoutDependency) {
-    if (timeoutDependencies == null)
-      timeoutDependencies = new ArrayList<>(3);
-    timeoutDependencies.add((Future) timeoutDependency);
+  synchronized void injectTimeout(Future<?> timeoutFuture) {
+    if (timeoutFutures == null)
+      timeoutFutures = new ArrayList<>(3);
+    timeoutFutures.add((Future) timeoutFuture);
   }
 }

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

@@ -96,7 +96,7 @@ protected Supplier<CompletableFuture<ExecutionResult>> supplyAsync(
           if (!policy.isAsync())
             callable.call();
           else
-            future.inject(scheduler.schedule(callable, result.getWaitNanos(), TimeUnit.NANOSECONDS));
+            future.injectPolicy(scheduler.schedule(callable, result.getWaitNanos(), TimeUnit.NANOSECONDS));
         } catch (Throwable t) {
           promise.completeExceptionally(t);
         }

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

@@ -106,7 +106,7 @@ static Supplier<CompletableFuture<ExecutionResult>> getPromiseAsync(
 
         try {
           scheduled.set(true);
-          future.inject(scheduler.schedule(callable, 0, TimeUnit.NANOSECONDS));
+          future.injectPolicy(scheduler.schedule(callable, 0, TimeUnit.NANOSECONDS));
         } catch (Throwable t) {
           promise.completeExceptionally(t);
         }
@@ -143,14 +143,18 @@ public synchronized CompletableFuture<ExecutionResult> get() {
    * Returns a Supplier that pre-executes the {@code execution}, applies the {@code supplier}, records the result and
    * returns a promise containing the result.
    */
+  @SuppressWarnings("rawtypes")
   static <T> Supplier<CompletableFuture<ExecutionResult>> getPromiseOfStage(
-    ContextualSupplier<? extends CompletionStage<? extends T>> supplier, AbstractExecution execution) {
+    ContextualSupplier<? extends CompletionStage<? extends T>> supplier, AsyncExecution execution) {
     Assert.notNull(supplier, "supplier");
     return () -> {
       CompletableFuture<ExecutionResult> promise = new CompletableFuture<>();
       try {
         execution.preExecute();
-        supplier.get(execution).whenComplete((result, failure) -> {
+        CompletionStage<? extends T> stage = supplier.get(execution);
+        if (stage instanceof Future)
+          execution.future.injectPolicy((Future) stage);
+        stage.whenComplete((result, failure) -> {
           if (failure instanceof CompletionException)
             failure = failure.getCause();
           ExecutionResult r = failure == null ? ExecutionResult.success(result) : ExecutionResult.failure(failure);

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

@@ -133,7 +133,7 @@ else if (postResult != null) {
                               retryScheduledListener.handle(postResult, execution);
 
                             previousResult = postResult;
-                            future.inject(scheduler.schedule(this, postResult.getWaitNanos(), TimeUnit.NANOSECONDS));
+                            future.injectPolicy(scheduler.schedule(this, postResult.getWaitNanos(), TimeUnit.NANOSECONDS));
                             future.injectCancelFn(() -> {
                               // Ensure that the promise completes if a scheduled retry is cancelled
                               if (executionCancelled())

src/test/java/net/jodah/failsafe/issues/Issue266Test.java

@@ -0,0 +1,42 @@
+package net.jodah.failsafe.issues;
+
+import net.jodah.failsafe.Asserts;
+import net.jodah.failsafe.Failsafe;
+import net.jodah.failsafe.Policy;
+import net.jodah.failsafe.RetryPolicy;
+import org.testng.annotations.Test;
+
+import java.util.concurrent.CancellationException;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionStage;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.testng.Assert.assertEquals;
+
+@Test
+public class Issue266Test {
+  AtomicInteger cancelledCounter = new AtomicInteger();
+
+  public void test() {
+    Policy<String> retryPolicy = new RetryPolicy<>();
+    CompletableFuture<String> future = Failsafe.with(retryPolicy).getStageAsync(this::computeSomething);
+    future.whenComplete((r, t) -> {
+      if (t instanceof CancellationException) {
+        cancelledCounter.incrementAndGet();
+      }
+    });
+    future.cancel(true);
+    Asserts.assertThrows(future::get, CancellationException.class);
+    assertEquals(cancelledCounter.get(), 2);
+  }
+
+  CompletionStage<String> computeSomething() {
+    CompletableFuture<String> future = new CompletableFuture<>();
+    future.whenComplete((r, t) -> {
+      if (t instanceof CancellationException) {
+        cancelledCounter.incrementAndGet();
+      }
+    });
+    return future;
+  }
+}