Added support for Executor with getStageAsync

Author: jhalterman

src/main/java/dev/failsafe/FailsafeExecutor.java

@@ -210,7 +210,7 @@ public <T extends R> CompletableFuture<T> getAsyncExecution(AsyncRunnable<T> run
    * @throws RejectedExecutionException if the {@code supplier} cannot be scheduled for execution
    */
   public <T extends R> CompletableFuture<T> getStageAsync(CheckedSupplier<? extends CompletionStage<T>> supplier) {
-    return callAsync(future -> getPromiseOfStage(toCtxSupplier(supplier), future), false);
+    return callAsync(future -> getPromiseOfStage(toCtxSupplier(supplier), future, executor), false);
   }
 
   /**
@@ -219,6 +219,7 @@ public <T extends R> CompletableFuture<T> getStageAsync(CheckedSupplier<? extend
    * <p>Cancelling the resulting {@link CompletableFuture} will automatically cancels the supplied {@link
    * CompletionStage} if it's a {@link Future}.</p>
    * <ul>
+   *   <li>If the {@code supplier} returns {@code null}, the execution attempt will record a {@code null} result.</li>
    *   <li>If the execution fails because a {@link Timeout} is exceeded, the resulting future is completed exceptionally
    *   with {@link TimeoutExceededException}.</li>
    *   <li>If the execution fails because a {@link CircuitBreaker} is open, the resulting future is completed
@@ -232,7 +233,7 @@ public <T extends R> CompletableFuture<T> getStageAsync(CheckedSupplier<? extend
    */
   public <T extends R> CompletableFuture<T> getStageAsync(
     ContextualSupplier<T, ? extends CompletionStage<T>> supplier) {
-    return callAsync(future -> getPromiseOfStage(supplier, future), false);
+    return callAsync(future -> getPromiseOfStage(supplier, future, executor), false);
   }
 
   /**

src/main/java/dev/failsafe/Functions.java

@@ -125,26 +125,33 @@ public synchronized CompletableFuture<ExecutionResult<R>> apply(AsyncExecutionIn
    */
   @SuppressWarnings("unchecked")
   static <R> Function<AsyncExecutionInternal<R>, CompletableFuture<ExecutionResult<R>>> getPromiseOfStage(
-    ContextualSupplier<R, ? extends CompletionStage<? extends R>> supplier, FailsafeFuture<R> future) {
+    ContextualSupplier<R, ? extends CompletionStage<? extends R>> supplier, FailsafeFuture<R> future,
+    Executor executor) {
 
     Assert.notNull(supplier, "supplier");
     return execution -> {
       CompletableFuture<ExecutionResult<R>> promise = new CompletableFuture<>();
       try {
         execution.preExecute();
-        CompletionStage<? extends R> stage = supplier.get(execution);
-
-        // Propagate outer cancellations to the stage
-        if (stage instanceof Future)
-          future.propagateCancellation((Future<R>) stage);
+        CompletionStage<? extends R> stage = withExecutor(supplier, executor).get(execution);
 
-        stage.whenComplete((result, failure) -> {
-          if (failure instanceof CompletionException)
-            failure = failure.getCause();
-          ExecutionResult<R> r = failure == null ? ExecutionResult.success(result) : ExecutionResult.failure(failure);
+        if (stage == null) {
+          ExecutionResult<R> r = ExecutionResult.success(null);
           execution.record(r);
           promise.complete(r);
-        });
+        } else {
+          // Propagate outer cancellations to the stage
+          if (stage instanceof Future)
+            future.propagateCancellation((Future<R>) stage);
+
+          stage.whenComplete((result, failure) -> {
+            if (failure instanceof CompletionException)
+              failure = failure.getCause();
+            ExecutionResult<R> r = failure == null ? ExecutionResult.success(result) : ExecutionResult.failure(failure);
+            execution.record(r);
+            promise.complete(r);
+          });
+        }
       } catch (Throwable t) {
         ExecutionResult<R> result = ExecutionResult.failure(t);
         execution.record(result);

src/test/java/dev/failsafe/issues/Issue311Test.java

@@ -2,25 +2,57 @@
 
 import dev.failsafe.Failsafe;
 import dev.failsafe.RetryPolicy;
+import net.jodah.concurrentunit.Waiter;
 import org.testng.annotations.Test;
 
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Executor;
 import java.util.concurrent.Executors;
 import java.util.concurrent.atomic.AtomicInteger;
 
-import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.*;
 
+@Test
 public class Issue311Test {
-  @Test
-  void failsafeFail() throws Throwable {
+  public void failsafeFail() throws Throwable {
     AtomicInteger counter = new AtomicInteger(0);
     Executor executor = Executors.newSingleThreadExecutor();
     Failsafe.with(RetryPolicy.builder().handle(RuntimeException.class).withMaxAttempts(2).build())
       .with(executor)
       .runAsync(() -> {
         if (counter.incrementAndGet() == 1)
           throw new RuntimeException();
-      }).get();
+      })
+      .get();
     assertEquals(counter.get(), 2);
   }
+
+  public void testNullCompletionStage() throws Throwable {
+    assertNull(Failsafe.none().getStageAsync(() -> {
+      return null;
+    }).get());
+  }
+
+  public void testRunAsyncWithThreadLocalInExecutor() throws Throwable {
+    ThreadLocal<Boolean> threadLocal = new ThreadLocal<>();
+    Executor executor = runnable -> {
+      threadLocal.set(true);
+      runnable.run();
+    };
+    Failsafe.none().with(executor).runAsync(() -> {
+      assertTrue(threadLocal.get());
+    }).get();
+  }
+
+  public void testGetStageAsyncWithThreadLocalInExecutor() throws Throwable {
+    ThreadLocal<Boolean> threadLocal = new ThreadLocal<>();
+    Executor executor = runnable -> {
+      threadLocal.set(true);
+      runnable.run();
+    };
+    assertNull(Failsafe.none().with(executor).getStageAsync(() -> {
+      assertTrue(threadLocal.get());
+      return CompletableFuture.completedFuture("ignored");
+    }).get());
+  }
 }