new: allow cancellations in upgrade and downgrade process

RelativityMC · Nov 9, 2024 · 1bd210e · 1bd210e
1 parent ac9b234
commit 1bd210e
Show file tree

Hide file tree

Showing 5 changed files with 64 additions and 15 deletions.
diff --git a/src/main/java/com/ishland/flowsched/scheduler/Cancellable.java b/src/main/java/com/ishland/flowsched/scheduler/Cancellable.java
@@ -0,0 +1,17 @@
+package com.ishland.flowsched.scheduler;
+
+import java.util.concurrent.atomic.AtomicBoolean;
+
+public class Cancellable {
+
+    private final AtomicBoolean cancelled = new AtomicBoolean(false);
+
+    public void cancel() {
+        this.cancelled.set(true);
+    }
+
+    public boolean isCancelled() {
+        return this.cancelled.get();
+    }
+
+}
diff --git a/src/main/java/com/ishland/flowsched/scheduler/ItemHolder.java b/src/main/java/com/ishland/flowsched/scheduler/ItemHolder.java
@@ -190,7 +190,7 @@ public void consolidateMarkDirty(StatusAdvancingScheduler<K, V, Ctx, ?> schedule
         this.busyRefCounter.addListenerOnce(() -> scheduler.markDirty(this.getKey()));
     }
 
-    public boolean setStatus(ItemStatus<K, V, Ctx> status) {
+    public boolean setStatus(ItemStatus<K, V, Ctx> status, boolean isCancellation) {
         assertOpen();
         ItemTicket<K, V, Ctx>[] ticketsToFire = null;
         CompletableFuture<Void> futureToFire = null;
@@ -224,8 +224,10 @@ public boolean setStatus(ItemStatus<K, V, Ctx> status) {
                 this.status = status;
                 final CompletableFuture<Void> future = this.futures[status.ordinal()];
 
-                Assertions.assertTrue(future != UNLOADED_FUTURE);
-                Assertions.assertTrue(!future.isDone());
+                if (!isCancellation) {
+                    Assertions.assertTrue(future != UNLOADED_FUTURE);
+                    Assertions.assertTrue(!future.isDone());
+                }
                 futureToFire = future;
                 ticketsToFire = this.tickets.getTicketsForStatus(status).toArray(ItemTicket[]::new);
             }

diff --git a/src/main/java/com/ishland/flowsched/scheduler/ItemStatus.java b/src/main/java/com/ishland/flowsched/scheduler/ItemStatus.java
@@ -37,9 +37,9 @@ default ItemStatus<K, V, Ctx> getNext() {
 
     int ordinal();
 
-    CompletionStage<Void> upgradeToThis(Ctx context);
+    CompletionStage<Void> upgradeToThis(Ctx context, Cancellable cancellable);
 
-    CompletionStage<Void> downgradeFromThis(Ctx context);
+    CompletionStage<Void> downgradeFromThis(Ctx context, Cancellable cancellable);
 
     /**
      * Get the dependencies of the given item at the given status.

diff --git a/src/main/java/com/ishland/flowsched/scheduler/StatusAdvancingScheduler.java b/src/main/java/com/ishland/flowsched/scheduler/StatusAdvancingScheduler.java
@@ -170,7 +170,7 @@ public boolean tick() {
                     }
                     holder.submitOp(CompletableFuture.runAsync(() -> advanceStatus0(holder, nextStatus, key), getBackgroundExecutor()));
                 } else {
-                    final boolean success = holder.setStatus(nextStatus);
+                    final boolean success = holder.setStatus(nextStatus, false);
                     if (!success) {
                         continue; // target status is modified to be higher
                     }
@@ -186,10 +186,12 @@ private void downgradeStatus0(ItemHolder<K, V, Ctx, UserData> holder, ItemStatus
         final KeyStatusPair<K, V, Ctx>[] dependencies = holder.getDependencies(current);
         Assertions.assertTrue(dependencies != null, "No dependencies for downgrade");
 
+        Cancellable cancellable = new Cancellable();
+
         final Completable completable = Completable.defer(() -> {
                     Assertions.assertTrue(holder.isBusy());
                     final Ctx ctx = makeContext(holder, current, dependencies, false);
-                    final CompletionStage<Void> stage = current.downgradeFromThis(ctx);
+                    final CompletionStage<Void> stage = current.downgradeFromThis(ctx, cancellable);
                     return Completable.fromCompletionStage(stage);
                 })
                 .subscribeOn(getSchedulerBackedByBackgroundExecutor())
@@ -198,6 +200,16 @@ private void downgradeStatus0(ItemHolder<K, V, Ctx, UserData> holder, ItemStatus
                     try {
                         Assertions.assertTrue(holder.isBusy());
 
+                        {
+                            Throwable actual = throwable;
+                            while (actual instanceof CompletionException ex) actual = ex.getCause();
+                            if (cancellable.isCancelled() && actual instanceof CancellationException) {
+                                holder.setStatus(current, true);
+                                holder.consolidateMarkDirty(this);
+                                return;
+                            }
+                        }
+
                         final ExceptionHandlingAction action = this.tryHandleTransactionException(holder, nextStatus, false, throwable);
                         switch (action) {
                             case PROCEED -> {
@@ -221,10 +233,10 @@ private void advanceStatus0(ItemHolder<K, V, Ctx, UserData> holder, ItemStatus<K
         // Advance
         final KeyStatusPair<K, V, Ctx>[] dependencies = nextStatus.getDependencies(holder);
         final CancellationSignaller dependencyCompletable = getDependencyFuture0(dependencies, holder, nextStatus);
-        AtomicBoolean isCancelled = new AtomicBoolean(false);
+        Cancellable cancellable = new Cancellable();
 
         CancellationSignaller signaller = new CancellationSignaller(unused -> {
-            isCancelled.set(true);
+            cancellable.cancel();
             dependencyCompletable.cancel();
         });
 
@@ -239,7 +251,7 @@ private void advanceStatus0(ItemHolder<K, V, Ctx, UserData> holder, ItemStatus<K
                 .andThen(Completable.defer(() -> {
                     Assertions.assertTrue(holder.isBusy());
                     final Ctx ctx = makeContext(holder, nextStatus, dependencies, false);
-                    final CompletionStage<Void> stage = nextStatus.upgradeToThis(ctx);
+                    final CompletionStage<Void> stage = nextStatus.upgradeToThis(ctx, cancellable);
                     return Completable.fromCompletionStage(stage).cache();
                 }))
                 .observeOn(getSchedulerBackedByBackgroundExecutor())
@@ -250,7 +262,15 @@ private void advanceStatus0(ItemHolder<K, V, Ctx, UserData> holder, ItemStatus<K
                         {
                             Throwable actual = throwable;
                             while (actual instanceof CompletionException ex) actual = ex.getCause();
-                            if (isCancelled.get() && actual instanceof CancellationException) {
+                            if (cancellable.isCancelled() && actual instanceof CancellationException) {
+                                if (holder.getDependencies(nextStatus) != null) {
+                                    releaseDependencies(holder, nextStatus);
+                                }
+                                try {
+                                    signaller.fireComplete(actual);
+                                } catch (Throwable t) {
+                                    t.printStackTrace();
+                                }
                                 holder.consolidateMarkDirty(this);
                                 return;
                             }
@@ -261,7 +281,7 @@ private void advanceStatus0(ItemHolder<K, V, Ctx, UserData> holder, ItemStatus<K
                         final ExceptionHandlingAction action = this.tryHandleTransactionException(holder, nextStatus, true, throwable);
                         switch (action) {
                             case PROCEED -> {
-                                holder.setStatus(nextStatus);
+                                holder.setStatus(nextStatus, false);
                                 rerequestDependencies(holder, nextStatus);
                                 holder.consolidateMarkDirty(this);
                                 onItemUpgrade(holder, nextStatus);

diff --git a/src/test/java/com/ishland/flowsched/scheduler/support/TestStatus.java b/src/test/java/com/ishland/flowsched/scheduler/support/TestStatus.java
@@ -1,12 +1,14 @@
 package com.ishland.flowsched.scheduler.support;
 
+import com.ishland.flowsched.scheduler.Cancellable;
 import com.ishland.flowsched.scheduler.ItemHolder;
 import com.ishland.flowsched.scheduler.ItemStatus;
 import com.ishland.flowsched.scheduler.KeyStatusPair;
 
 import java.util.ArrayList;
-import java.util.Collection;
 import java.util.List;
+import java.util.Random;
+import java.util.concurrent.CancellationException;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
 
@@ -30,14 +32,22 @@ public ItemStatus<Long, TestItem, TestContext>[] getAllStatuses() {
     }
 
     @Override
-    public CompletionStage<Void> upgradeToThis(TestContext context) {
+    public CompletionStage<Void> upgradeToThis(TestContext context, Cancellable cancellable) {
 //        System.out.println(String.format("Upgrading %d to %s", context.key(), this));
+        if (new Random().nextBoolean()) {
+            cancellable.cancel();
+            return CompletableFuture.failedFuture(new CancellationException());
+        }
         return CompletableFuture.completedFuture(null);
     }
 
     @Override
-    public CompletionStage<Void> downgradeFromThis(TestContext context) {
+    public CompletionStage<Void> downgradeFromThis(TestContext context, Cancellable cancellable) {
 //        System.out.println(String.format("Downgrading %d from %s", context.key(), this));
+        if (new Random().nextBoolean()) {
+            cancellable.cancel();
+            return CompletableFuture.failedFuture(new CancellationException());
+        }
         return CompletableFuture.completedFuture(null);
     }