Fix currentStage cancelation in AsyncLoop; optimize intermediate stage cancelation in AbstractCompletableTask (thenCompose / exceptionallyCompose)

Author: vsilaev

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

@@ -28,7 +28,6 @@
 import static net.tascalate.concurrent.SharedFunctions.wrapCompletionException;
 import static net.tascalate.concurrent.SharedFunctions.wrapExecutionException;
 
-import java.util.Arrays;
 import java.util.concurrent.Callable;
 import java.util.concurrent.CancellationException;
 import java.util.concurrent.CompletableFuture;
@@ -66,31 +65,18 @@ protected AbstractCompletableTask(Executor defaultExecutor, Callable<T> action)
         this.task = new StageTransition(action);
     }
 
-    private CompletionStage<?>[] cancellableOrigins;
-    private Object cancellableOriginsLock = new Object();
-    
-    protected void resetCancellableOrigins(CompletionStage<?>... origins) {
-        synchronized (cancellableOriginsLock) {
-            this.cancellableOrigins = origins; 
-        }
-    }
-    
-    protected void cancelOrigins(boolean mayInterruptIfRunning) {
-        synchronized (cancellableOriginsLock) {
-            if (null == cancellableOrigins) {
-                return;
-            }
-            Arrays.stream(cancellableOrigins).forEach(p -> cancelPromise(p, mayInterruptIfRunning)); 
-        }
-    }
+    private volatile CompletionStage<?> intermediateStage;
 
     abstract void fireTransition(Callable<T> code);
 
     @Override
     public boolean cancel(boolean mayInterruptIfRunning) {
         if (task.cancel(mayInterruptIfRunning)) {
             failure(new CancellationException());
-            cancelOrigins(mayInterruptIfRunning);
+            CompletionStage<?> s = intermediateStage;
+            if (null != s) {
+                cancelPromise(s, mayInterruptIfRunning);
+            }
             return true;
         } else {
             return false;
@@ -242,7 +228,7 @@ public <U> Promise<U> thenComposeAsync(Function<? super T, ? extends CompletionS
         AbstractCompletableTask<U> nextStage = internalCreateCompletionStage(executor);
         // Need to enlist tempStage while it is non-visible outside
         // and may not be used to interrupt fn.apply();
-        nextStage.resetCancellableOrigins(tempStage);
+        nextStage.intermediateStage = tempStage;
 
         // We must ALWAYS run through the execution
         // of nextStage.task when this nextStage is
@@ -257,10 +243,16 @@ public <U> Promise<U> thenComposeAsync(Function<? super T, ? extends CompletionS
             tempStage, 
             consumerAsFunction(r -> {
                 try {
-                    CompletionStage<U> returned = fn.apply(r);
                     // tempStage is completed successfully, so no sense
                     // to include it in cancellableOrigins
-                    // However, nextStage is in progress
+                    nextStage.intermediateStage = null;
+                    if (nextStage.isDone()) {
+                        // Was canceled by user code
+                        // Don't need to run anything
+                        return;
+                    }
+                    CompletionStage<U> returned = fn.apply(r);
+                    // nextStage is in progress
                     // IMPORTANT: it COULD be shared, but typically is not
                     // So in very rare case some nasty behavior MAY exist 
                     // if others depends on it
@@ -280,16 +272,15 @@ public <U> Promise<U> thenComposeAsync(Function<? super T, ? extends CompletionS
                     //} catch (InterruptedException ex) {
                     //}
 
-                    nextStage.resetCancellableOrigins(returned);
+                    // Always assign before check for cancellation to avoid race
+                    nextStage.intermediateStage = returned;
                     if (nextStage.isCancelled()) {
-                        nextStage.cancelOrigins(true);
+                        cancelPromise(returned, true);
                     } else {
                         // Synchronous, while transition to tempStage is asynchronous already
                         returned.whenComplete(biConsumer(onResult, onError));
                     }
                 } catch (Throwable e) {
-                    // must-have if fn.apply above failed 
-                    nextStage.resetCancellableOrigins((CompletionStage<U>)null);
                     // no need to check nextStage.isCancelled()
                     // while there are no origins to cancel
                     // propagate error immediately
@@ -320,7 +311,7 @@ public Promise<T> exceptionallyComposeAsync(Function<Throwable, ? extends Comple
         AbstractCompletableTask<Void> tempStage = internalCreateCompletionStage(executor);
         AbstractCompletableTask<T> nextStage = internalCreateCompletionStage(executor);
 
-        nextStage.resetCancellableOrigins(tempStage);
+        nextStage.intermediateStage = tempStage;
 
         Consumer<? super T> onResult = nextStage.runTransition(Function.identity());
         Consumer<? super Throwable> onError = nextStage.runTransition(AbstractCompletableTask::forwardException);
@@ -330,15 +321,20 @@ public Promise<T> exceptionallyComposeAsync(Function<Throwable, ? extends Comple
             consumerAsFunction(onResult),
             consumerAsFunction(error -> {
                 try {
+                    nextStage.intermediateStage = null;
+                    if (nextStage.isDone()) {
+                        // Was canceled by user code
+                        // Don't need to run anything
+                        return;
+                    }
                     CompletionStage<T> returned = fn.apply(error);
-                    nextStage.resetCancellableOrigins(returned);
+                    nextStage.intermediateStage = returned;
                     if (nextStage.isCancelled()) {
-                        nextStage.cancelOrigins(true);
+                        cancelPromise(returned, true);
                     } else {
                         returned.whenComplete(biConsumer(onResult, onError));
                     }
                 } catch (Throwable e) {
-                    nextStage.resetCancellableOrigins((CompletionStage<T>)null);
                     // In JDK 12 CompletionStage.composeExceptionally[Async] uses *.handle[Async]
                     // So overwrite returned error with the latest one - as in handle()
                     e.addSuppressed(error);

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

@@ -35,7 +35,10 @@ class AsyncLoop<T> extends CompletableFutureWrapper<T> {
     @Override
     public boolean cancel(boolean mayInterruptIfRunning) {
         if (super.cancel(mayInterruptIfRunning)) {
-            cancelCurrentStage(mayInterruptIfRunning);
+            CompletionStage<?> s = currentStage;
+            if (null != s) {
+                SharedFunctions.cancelPromise(s, mayInterruptIfRunning);
+            }  
             return true;
         } else {
             return false;
@@ -60,17 +63,25 @@ private void run(T resolvedValue, Thread initialThread, IterationState<T> initia
                     if (isDone()) {
                         break;
                     } else if (loopCondition.test(currentValue)) {
-                        (currentStage = loopBody.apply(currentValue)).whenComplete((next, ex) -> {
-                            if (ex != null) {
-                                failure(ex);
-                            } else {
-                                run(next, currentThread, currentState);
-                            }
-                        });
+                        CompletionStage<T> returned = loopBody.apply(currentValue);
+                        // Assign before check to avoid race
+                        currentStage = returned;
+                        // isDone() is never slower than isCancel() -- 
+                        // actually, the test is for cancellation
+                        // but isDone() is ok.
                         if (isDone()) {
                             // If race between this.cancel and this.run
                             // Double-cancel is not an issue
-                            cancelCurrentStage(true);
+                            SharedFunctions.cancelPromise(returned, true);
+                            break;
+                        } else {
+                            returned.whenComplete((next, ex) -> {
+                                if (ex != null) {
+                                    failure(ex);
+                                } else {
+                                    run(next, currentThread, currentState);
+                                }
+                            });
                         }
                     } else {
                         success(currentValue);
@@ -86,13 +97,6 @@ private void run(T resolvedValue, Thread initialThread, IterationState<T> initia
         }
     }
 
-    private void cancelCurrentStage(boolean mayInterruptIfRunning) {
-        CompletionStage<?> dependent = currentStage;
-        if (null != dependent) {
-            SharedFunctions.cancelPromise(dependent, mayInterruptIfRunning);
-        }        
-    }
-    
     private static final class IterationState<T> {
         static final Object END = new Object();
 

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

@@ -87,6 +87,13 @@ protected ConfigurableDependentPromise(Promise<T> delegate,
         this.cancellableOrigins = cancellableOrigins;
     }
 
+    protected ConfigurableDependentPromise<T> setup() {
+        if (!isEmptyArray(cancellableOrigins)) {
+            delegate.onCancel(() -> cancelPromises(cancellableOrigins, true));
+        }
+        return this;
+    }
+    
     public static <U> DependentPromise<U> from(Promise<U> source) {
         return from(source, PromiseOrigin.NONE);
     }
@@ -102,8 +109,7 @@ protected <U> DependentPromise<U> wrap(Promise<U> original, CompletionStage<?>[]
     private static <U> DependentPromise<U> doWrap(Promise<U> original, 
                                                   Set<PromiseOrigin> defaultEnlistOptions, 
                                                   CompletionStage<?>[] cancellableOrigins) {
-        boolean noOrigins = isEmptyArray(cancellableOrigins); 
-        if (noOrigins) {
+        if (isEmptyArray(cancellableOrigins)) {
             // Nothing to enlist additionally for this "original" instance
             if (original instanceof ConfigurableDependentPromise) {
                 ConfigurableDependentPromise<U> ioriginal = (ConfigurableDependentPromise<U>)original;
@@ -114,10 +120,7 @@ private static <U> DependentPromise<U> doWrap(Promise<U> original,
             }
         }
 
-        return new ConfigurableDependentPromise<>(
-            noOrigins ? original : original.onCancel(() -> cancelPromises(cancellableOrigins, true)), 
-            defaultEnlistOptions, cancellableOrigins
-        );
+        return new ConfigurableDependentPromise<>(original, defaultEnlistOptions, cancellableOrigins).setup();
     }
 
     @Override

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

@@ -100,9 +100,11 @@ default boolean isCompletedExceptionally() {
     }
 
     default Promise<T> onCancel(Runnable action) {
-        if (isCancelled()) {
-            action.run();
-        } else if (!isDone()) {
+        if (isDone()) {
+            if (isCancelled()) {
+                action.run();
+            }            
+        } else {
             exceptionally(__ -> {
                if (isCancelled()) {
                    action.run();