Ensure CompletedStage waits for other dependents

Previously a CompletedStage which held an exception would immediately
complete the returned stage in methods like thenCombine with this
exception. Instead it should wait for the other dependent stage to
complete before completing the returned stage. SimpleCompletionStage was
found not to have this incorrect behavior. Fixes #49

Author: Renar Narubin

asyncutil/src/main/java/com/ibm/asyncutil/util/CompletedStage.java

@@ -178,87 +178,111 @@ public CompletionStage<Void> thenRunAsync(final Runnable action, final Executor
   public <U, V> CompletionStage<V> thenCombine(final CompletionStage<? extends U> other,
       final BiFunction<? super T, ? super U, ? extends V> fn) {
     Objects.requireNonNull(fn);
-    if (this.exception == null) {
-      return other.thenApply(u -> fn.apply(this.result, u));
-    }
-    return typedException();
+    return other.thenApply(
+        this.exception == null
+            ? u -> fn.apply(this.result, u)
+            : u -> {
+              throw CompletedStage.wrapIfNecessary(this.exception);
+            });
   }
 
   @Override
   public <U, V> CompletionStage<V> thenCombineAsync(final CompletionStage<? extends U> other,
       final BiFunction<? super T, ? super U, ? extends V> fn) {
     Objects.requireNonNull(fn);
-    if (this.exception == null) {
-      return other.thenApplyAsync(u -> fn.apply(this.result, u));
-    }
-    return typedException();
+    return other.thenApplyAsync(
+        this.exception == null
+            ? u -> fn.apply(this.result, u)
+            : u -> {
+              throw CompletedStage.wrapIfNecessary(this.exception);
+            });
   }
 
   @Override
   public <U, V> CompletionStage<V> thenCombineAsync(final CompletionStage<? extends U> other,
       final BiFunction<? super T, ? super U, ? extends V> fn, final Executor executor) {
     Objects.requireNonNull(fn);
-    if (this.exception == null) {
-      return other.thenApplyAsync(u -> fn.apply(this.result, u), executor);
-    }
-    return typedException();
+    return other.thenApplyAsync(
+        this.exception == null
+            ? u -> fn.apply(this.result, u)
+            : u -> {
+              throw CompletedStage.wrapIfNecessary(this.exception);
+            },
+        executor);
   }
 
   @Override
   public <U> CompletionStage<Void> thenAcceptBoth(final CompletionStage<? extends U> other,
       final BiConsumer<? super T, ? super U> action) {
     Objects.requireNonNull(action);
-    if (this.exception == null) {
-      return other.thenAccept(u -> action.accept(this.result, u));
-    }
-    return typedException();
+    return other.thenAccept(
+        this.exception == null
+            ? u -> action.accept(this.result, u)
+            : u -> {
+              throw CompletedStage.wrapIfNecessary(this.exception);
+            });
   }
 
   @Override
   public <U> CompletionStage<Void> thenAcceptBothAsync(final CompletionStage<? extends U> other,
       final BiConsumer<? super T, ? super U> action) {
     Objects.requireNonNull(action);
-    if (this.exception == null) {
-      return other.thenAcceptAsync(u -> action.accept(this.result, u));
-    }
-    return typedException();
+    return other.thenAcceptAsync(
+        this.exception == null
+            ? u -> action.accept(this.result, u)
+            : u -> {
+              throw CompletedStage.wrapIfNecessary(this.exception);
+            });
   }
 
   @Override
   public <U> CompletionStage<Void> thenAcceptBothAsync(final CompletionStage<? extends U> other,
       final BiConsumer<? super T, ? super U> action, final Executor executor) {
     Objects.requireNonNull(action);
-    if (this.exception == null) {
-      return other.thenAcceptAsync(u -> action.accept(this.result, u), executor);
-    }
-    return typedException();
+    return other.thenAcceptAsync(
+        this.exception == null
+            ? u -> action.accept(this.result, u)
+            : u -> {
+              throw CompletedStage.wrapIfNecessary(this.exception);
+            },
+        executor);
   }
 
   @Override
   public CompletionStage<Void> runAfterBoth(final CompletionStage<?> other, final Runnable action) {
-    if (this.exception == null) {
-      return other.thenRun(action);
-    }
-    return typedException();
+    Objects.requireNonNull(action);
+    return other.thenRun(
+        this.exception == null
+            ? action
+            : () -> {
+              throw CompletedStage.wrapIfNecessary(this.exception);
+            });
   }
 
   @Override
   public CompletionStage<Void> runAfterBothAsync(final CompletionStage<?> other,
       final Runnable action) {
-    if (this.exception == null) {
-      return other.thenRunAsync(action);
-    }
-    return typedException();
+    Objects.requireNonNull(action);
+    return other.thenRunAsync(
+        this.exception == null
+            ? action
+            : () -> {
+              throw CompletedStage.wrapIfNecessary(this.exception);
+            });
   }
 
   @Override
   public CompletionStage<Void> runAfterBothAsync(final CompletionStage<?> other,
       final Runnable action,
       final Executor executor) {
-    if (this.exception == null) {
-      return other.thenRunAsync(action, executor);
-    }
-    return typedException();
+    Objects.requireNonNull(action);
+    return other.thenRunAsync(
+        this.exception == null
+            ? action
+            : () -> {
+              throw CompletedStage.wrapIfNecessary(this.exception);
+            },
+        executor);
   }
 
   @Override

asyncutil/src/test/java/com/ibm/asyncutil/util/CombinatorsTest.java

@@ -11,14 +11,19 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
 import java.util.concurrent.CompletionStage;
+import java.util.concurrent.Executor;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import java.util.function.BiConsumer;
+import java.util.function.BiFunction;
 import java.util.function.Function;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
+import java.util.stream.Stream;
 
 import org.junit.Assert;
 import org.junit.Test;
@@ -31,6 +36,8 @@
 
 @RunWith(Parameterized.class)
 public class CombinatorsTest {
+  private static final Executor ASYNC = CompletableFuture::runAsync;
+
   private static class TestException extends RuntimeException {
     private static final long serialVersionUID = 1L;
   }
@@ -155,9 +162,9 @@ public void testAllOfError() {
         Arrays.asList(
             getCompletedStage(1),
             getExceptionalStage(new TestException()));
-    assertError(Combinators.allOf(futures));
-    assertError(Combinators.collect(futures));
-    assertError(Combinators.collect(futures, Collectors.toList()));
+    CombinatorsTest.assertError(Combinators.allOf(futures));
+    CombinatorsTest.assertError(Combinators.collect(futures));
+    CombinatorsTest.assertError(Combinators.collect(futures, Collectors.toList()));
   }
 
   @Test
@@ -173,15 +180,15 @@ public void testAllOfErrorNoShortCircuit() {
     final CompletionStage<List<Integer>> collCollect =
         Combinators.collect(futures, Collectors.toList());
 
-    assertIncomplete(voidAll);
-    assertIncomplete(collAll);
-    assertIncomplete(collCollect);
+    CombinatorsTest.assertIncomplete(voidAll);
+    CombinatorsTest.assertIncomplete(collAll);
+    CombinatorsTest.assertIncomplete(collCollect);
 
     delayed.complete(1);
 
-    assertError(voidAll);
-    assertError(collAll);
-    assertError(collCollect);
+    CombinatorsTest.assertError(voidAll);
+    CombinatorsTest.assertError(collAll);
+    CombinatorsTest.assertError(collCollect);
   }
 
   @Test
@@ -208,7 +215,7 @@ public void testKeyedAllError() {
               }
               return getCompletedStage(i);
             }));
-    assertError(Combinators.keyedAll(stageMap));
+    CombinatorsTest.assertError(Combinators.keyedAll(stageMap));
   }
 
   @Test
@@ -220,26 +227,84 @@ public void testKeyedAllErrorNoShortCircuit() {
     final CompletionStage<Map<Integer, Integer>> fut = Combinators.keyedAll(stageMap);
     int i = 0;
     for (final CompletableStage<Integer> future : stageMap.values()) {
-      assertIncomplete(fut);
+      CombinatorsTest.assertIncomplete(fut);
       if (i == 3) {
         future.completeExceptionally(new TestException());
       } else {
         future.complete(i);
       }
       i++;
     }
-    assertError(fut);
+    CombinatorsTest.assertError(fut);
+  }
+
+  /**
+   * Test that CompletionStage methods which depend on both of two stages always wait for both
+   * stages to complete.
+   */
+  @Test
+  public void testCombineEtAl() {
+    for (final BiFunction<CompletionStage<?>, CompletionStage<?>, CompletionStage<Void>> combineMethod : Stream
+        .<BiFunction<CompletionStage<?>, CompletionStage<?>, CompletionStage<Void>>>of(
+            (a, b) -> a.thenCombine(b, CombinatorsTest.voidFunction()),
+            (a, b) -> a.thenCombineAsync(b, CombinatorsTest.voidFunction()),
+            (a, b) -> a.thenCombineAsync(b, CombinatorsTest.voidFunction(), ASYNC),
+            (a, b) -> a.thenAcceptBoth(b, CombinatorsTest.voidConsumer()),
+            (a, b) -> a.thenAcceptBothAsync(b, CombinatorsTest.voidConsumer()),
+            (a, b) -> a.thenAcceptBothAsync(b, CombinatorsTest.voidConsumer(), ASYNC),
+            (a, b) -> a.runAfterBoth(b, CombinatorsTest.voidRunnable()),
+            (a, b) -> a.runAfterBothAsync(b, CombinatorsTest.voidRunnable()),
+            (a, b) -> a.runAfterBothAsync(b, CombinatorsTest.voidRunnable(), ASYNC))
+        // include all the functions in reverse as well, switch a and b in argument order
+        .flatMap(function -> Stream.of(function, (a, b) -> function.apply(b, a)))
+        .collect(Collectors.toList())) {
+      {
+        final CompletionStage<String> doneNormal = getCompletedStage("a");
+        final CompletableStage<String> incompleteNormal = getCompletableStage();
+
+        final CompletionStage<Void> combine = combineMethod.apply(doneNormal, incompleteNormal);
+
+        CombinatorsTest.assertIncomplete(combine);
+        incompleteNormal.complete("b");
+        TestUtil.join(combine);
+      }
+      {
+        final CompletionStage<String> doneExceptional = getExceptionalStage(new TestException());
+        final CompletableStage<String> incompleteNormal = getCompletableStage();
+
+        final CompletionStage<Void> combine =
+            combineMethod.apply(doneExceptional, incompleteNormal);
+
+        CombinatorsTest.assertIncomplete(combine);
+        incompleteNormal.complete("b");
+        CombinatorsTest.assertError(combine);
+      }
+    }
+  }
+
+  private static <T, V> BiFunction<T, V, Void> voidFunction() {
+    return (ig1, ig2) -> null;
+  }
+
+  private static <T, V> BiConsumer<T, V> voidConsumer() {
+    return (ig1, ig2) -> {
+    };
+  }
+
+  private static Runnable voidRunnable() {
+    return () -> {
+    };
   }
 
-  private <T> void assertError(final CompletionStage<T> stage) {
+  private static <T> void assertError(final CompletionStage<T> stage) {
     try {
       TestUtil.join(stage);
     } catch (final CompletionException e) {
       Assert.assertTrue(e.getCause() instanceof TestException);
     }
   }
 
-  private <T> void assertIncomplete(final CompletionStage<T> stage) {
+  private static <T> void assertIncomplete(final CompletionStage<T> stage) {
     try {
       TestUtil.join(stage, 20, TimeUnit.MILLISECONDS);
       Assert.fail("not all futures complete, get should timeout");