Adding DelayPolicy.withCustomizer to allow chaining calls for custom classes;

Making RetryPolicy.exceptionClassRetryable overridable;
Adding DependentPromise.as_ with more specific function arg type;
In AggregatingPromise cancel remaining before firing event;
Fix std. CompletableFuture.thenCompose[Async] behavior for cancellation -- now the Promise returned from fn is cancelled as well

Author: vsilaev

src/main/java/net/tascalate/concurrent/AggregatingPromise.java

@@ -42,6 +42,7 @@ abstract class AggregatingPromise<T, R> extends CompletablePromise<List<R>> {
     private final boolean cancelRemaining;
     private final List<? extends CompletionStage<? extends T>> promises;
 
+    @FunctionalInterface
     interface Constructor<T, R> {
         AggregatingPromise<T, R> create(int minResultsCount, int maxErrorsCount, boolean cancelRemaining,
                                         List<? extends CompletionStage<? extends T>> promises);
@@ -143,10 +144,14 @@ void onComplete(int idx, T result, Throwable error) {
                     // Synchronized around done
                     markRemainingCancelled();
                     // Now no other thread can modify results array.
-                    onSuccess(collectResults(minResultsCount, completions));
+                    
+                    // Cancel before firing events:
+                    // this helps with releasing thread pools, throttling etc.
                     if (cancelRemaining) {
                         cancelPromises();
                     }
+                    
+                    onSuccess(collectResults(minResultsCount, completions));
                 }
             }
         } else {
@@ -163,11 +168,14 @@ void onComplete(int idx, T result, Throwable error) {
                     // Synchronized around done
                     markRemainingCancelled();
                     // Now no other thread can modify errors array.
-                    onFailure(new MultitargetException(errors));
-
+                    
+                    // Cancel before firing events:
+                    // this helps with releasing thread pools, throttling etc.
                     if (cancelRemaining) {
                         cancelPromises();
                     }
+                    
+                    onFailure(new MultitargetException(errors));
                 }
             }
         }

src/main/java/net/tascalate/concurrent/CallbackRegistry.java

@@ -284,9 +284,9 @@ void callFailureCallback(Throwable failure) {
     }
 
     private static <S, U> void callCallback(Consumer<? super Callable<?>> stageTransition,
-                                    Function<? super S, ? extends U> callback, 
-                                    S value, 
-                                    Executor executor) {
+                                            Function<? super S, ? extends U> callback, 
+                                            S value,  
+                                            Executor executor) {
 
         Callable<U> callable = () -> callback.apply(value);
         try {

src/main/java/net/tascalate/concurrent/CompletableFutureWrapper.java

@@ -19,6 +19,8 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
 import java.util.concurrent.CompletionStage;
+import java.util.concurrent.Executor;
+import java.util.function.Function;
 
 import net.tascalate.concurrent.decorators.AbstractFutureDecorator;
 
@@ -30,7 +32,6 @@ protected CompletableFutureWrapper() {
         this(new CompletableFuture<T>());
     }
 
-    
     protected CompletableFutureWrapper(CompletableFuture<T> delegate) {
         super(delegate);
     }
@@ -49,4 +50,24 @@ public T join() throws CancellationException, CompletionException {
     protected <U> Promise<U> wrap(CompletionStage<U> original) {
         return new CompletableFutureWrapper<>((CompletableFuture<U>)original);
     }
+
+    // By default CompletableFuture doesn't interrupt a promise from thenCompose(fn)!
+    @Override
+    public <U> Promise<U> thenCompose(Function<? super T, ? extends CompletionStage<U>> fn) {
+        ComposedFutureRef<U> ref = new ComposedFutureRef<>();
+        return super.thenCompose(ref.captureResult(fn)).onCancel(ref.cancelCaptured);
+    }
+
+    @Override
+    public <U> Promise<U> thenComposeAsync(Function<? super T, ? extends CompletionStage<U>> fn) {
+        ComposedFutureRef<U> ref = new ComposedFutureRef<>();
+        return super.thenComposeAsync(ref.captureResult(fn)).onCancel(ref.cancelCaptured);
+    }
+
+    @Override
+    public <U> Promise<U> thenComposeAsync(Function<? super T, ? extends CompletionStage<U>> fn, 
+                                                    Executor executor) {
+        ComposedFutureRef<U> ref = new ComposedFutureRef<>();
+        return super.thenComposeAsync(ref.captureResult(fn), executor).onCancel(ref.cancelCaptured);        
+    }
 }

src/main/java/net/tascalate/concurrent/CompletionStageWrapper.java

@@ -25,11 +25,14 @@
 import java.util.concurrent.CompletionStage;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Executor;
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import java.util.function.Function;
 
 import net.tascalate.concurrent.decorators.AbstractCompletionStageDecorator;
+import net.tascalate.concurrent.decorators.AbstractPromiseDecorator;
 import net.tascalate.concurrent.decorators.BlockingCompletionStageDecorator;
 
 public class CompletionStageWrapper<T> 
@@ -117,20 +120,57 @@ public T join() {
             throw wrapCompletionException(fault);
         }        
     }
-    
 
     @Override
     protected <U> Promise<U> wrap(CompletionStage<U> original) {
         return Promises.from(original);
     }
 
     public static <T> Promise<T> from(CompletionStage<T> stage) {
+        return from(stage, true);
+    }
+    
+    public static <T> Promise<T> from(CompletionStage<T> stage, boolean composeSafe) {
+        Promise<T> result;
         if (stage instanceof Future) {
             // If we can delegate blocking Future API...
-            return BlockingCompletionStageDecorator.from(stage);
+            result = BlockingCompletionStageDecorator.from(stage);
         } else {
             // Otherwise fallback to own implementation
-            return new CompletionStageWrapper<>(stage);
+            result = new CompletionStageWrapper<>(stage);
         }
+        return composeSafe ? new ComposeSafePromise<>(result) : result;
+    }
+    
+    // By default CompletableFuture doesn't interrupt a promise from thenCompose(fn)!
+    // Pessimistically assume this "feature" for all CompletionStage impls
+    static class ComposeSafePromise<T> extends AbstractPromiseDecorator<T, Promise<T>> {
+        protected ComposeSafePromise(Promise<T> delegate) {
+            super(delegate);
+        }
+
+        @Override
+        protected <U> Promise<U> wrap(CompletionStage<U> original) {
+            return new ComposeSafePromise<>((Promise<U>)original);
+        }
+        
+        @Override
+        public <U> Promise<U> thenCompose(Function<? super T, ? extends CompletionStage<U>> fn) {
+            ComposedFutureRef<U> ref = new ComposedFutureRef<>();
+            return super.thenCompose(ref.captureResult(fn)).onCancel(ref.cancelCaptured);
+        }
+
+        @Override
+        public <U> Promise<U> thenComposeAsync(Function<? super T, ? extends CompletionStage<U>> fn) {
+            ComposedFutureRef<U> ref = new ComposedFutureRef<>();
+            return super.thenComposeAsync(ref.captureResult(fn)).onCancel(ref.cancelCaptured);
+        }
+
+        @Override
+        public <U> Promise<U> thenComposeAsync(Function<? super T, ? extends CompletionStage<U>> fn, 
+                                                        Executor executor) {
+            ComposedFutureRef<U> ref = new ComposedFutureRef<>();
+            return super.thenComposeAsync(ref.captureResult(fn), executor).onCancel(ref.cancelCaptured);        
+        }        
     }
 }

src/main/java/net/tascalate/concurrent/ComposedFutureRef.java

@@ -0,0 +1,39 @@
+/**
+ * Copyright 2015-2020 Valery Silaev (http://vsilaev.com)
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package net.tascalate.concurrent;
+
+import java.util.concurrent.CompletionStage;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Function;
+
+class ComposedFutureRef<U> extends AtomicReference<CompletionStage<U>> {
+    private static final long serialVersionUID = 1L;
+
+    <T> Function<? super T, ? extends CompletionStage<U>> captureResult(Function<? super T, ? extends CompletionStage<U>> fn) {
+        return v -> {
+            CompletionStage<U> result = fn.apply(v);
+            set(result);
+            return result;
+        };
+    }    
+
+    Runnable cancelCaptured = () -> {
+        CompletionStage<U> stage = get();
+        if (null != stage) {
+            SharedFunctions.cancelPromise(stage, true);
+        }
+    };
+}