fix: incorrect early ticket completion during downgrading

Author: ishland

src/main/java/com/ishland/flowsched/scheduler/Cancellable.java

@@ -1,17 +1,63 @@
 package com.ishland.flowsched.scheduler;
 
-import java.util.concurrent.atomic.AtomicBoolean;
+import com.ishland.flowsched.util.Assertions;
+
+import java.lang.invoke.MethodHandles;
+import java.lang.invoke.VarHandle;
 
 public class Cancellable {
 
-    private final AtomicBoolean cancelled = new AtomicBoolean(false);
+    @SuppressWarnings("unused")
+    private Runnable onCancel;
+
+    public boolean setup(Runnable onCancel) {
+        final var result = VH_CANCEL.compareAndExchangeAcquire(this, null, onCancel);
+        Assertions.assertTrue(result != COMPLETED, "Cancellation is already completed when setup");
+        return result == null;
+    }
+
+    public boolean complete() {
+        final var witness = (Runnable) VH_CANCEL.compareAndExchangeRelease(this, null, COMPLETED);
+        if (witness == CANCELLED || witness == COMPLETED) {
+            return false;
+        }
+        if (witness == null) {
+            return true;
+        }
+        // Set it to completed if we can; if we failed, it can only be cancel() or another complete()
+        return witness == VH_CANCEL.compareAndExchangeRelease(this, witness, COMPLETED);
+    }
 
-    public void cancel() {
-        this.cancelled.set(true);
+    public boolean cancel() {
+        final Runnable handle = (Runnable) VH_CANCEL.compareAndExchangeRelease(this, null, CANCELLED);
+        if (handle == null) {
+            return true;
+        }
+        if (handle != CANCELLED && handle != COMPLETED) {
+            handle.run();
+            // Set it to cancelled if we can; if we failed, it can only be complete() or another cancel()
+            return handle == VH_CANCEL.compareAndExchangeRelease(this, handle, CANCELLED);
+        }
+        return false;
     }
 
     public boolean isCancelled() {
-        return this.cancelled.get();
+        return VH_CANCEL.getAcquire(this) == CANCELLED;
     }
 
+    public boolean isCompleted() {
+        return VH_CANCEL.getAcquire(this) == COMPLETED;
+    }
+
+    private static final VarHandle VH_CANCEL;
+    private static final Runnable CANCELLED = () -> {};
+    private static final Runnable COMPLETED = () -> {};
+
+    static {
+        try {
+            VH_CANCEL = MethodHandles.lookup().findVarHandle(Cancellable.class, "onCancel", Runnable.class);
+        } catch (NoSuchFieldException | IllegalAccessException e) {
+            throw new RuntimeException(e);
+        }
+    }
 }

src/main/java/com/ishland/flowsched/scheduler/CancellationSignaller.java

@@ -54,7 +54,7 @@ public class ItemHolder<K, V, Ctx, UserData> {
     private final AtomicReference<V> item = new AtomicReference<>();
     private final AtomicReference<UserData> userData = new AtomicReference<>();
     private final BusyRefCounter busyRefCounter = new BusyRefCounter();
-    private final AtomicReference<Pair<CancellationSignaller, ItemStatus<K, V, Ctx>>> runningUpgradeAction = new AtomicReference<>();
+    private final AtomicReference<Pair<Cancellable, ItemStatus<K, V, Ctx>>> runningAction = new AtomicReference<>();
     private final TicketSet<K, V, Ctx> tickets;
     private volatile ItemStatus<K, V, Ctx> status = null;
 //    private final List<Pair<ItemStatus<K, V, Ctx>, Long>> statusHistory = ReferenceLists.synchronize(new ReferenceArrayList<>());
@@ -127,9 +127,9 @@ public synchronized boolean isBusy() {
         return busyRefCounter.isBusy();
     }
 
-    public ItemStatus<K, V, Ctx> upgradingStatusTo() {
+    public ItemStatus<K, V, Ctx> changingStatusTo() {
         assertOpen();
-        final Pair<CancellationSignaller, ItemStatus<K, V, Ctx>> pair = this.runningUpgradeAction.get();
+        final Pair<Cancellable, ItemStatus<K, V, Ctx>> pair = this.runningAction.get();
         return pair != null ? pair.right() : null;
     }
 
@@ -144,7 +144,9 @@ public void addTicket(ItemStatus<K, V, Ctx> targetStatus, ItemTicket ticket) {
             }
             this.tickets.addUnchecked(targetStatus);
             createFutures();
-            needConsumption = targetStatus.ordinal() <= this.getStatus().ordinal();
+            final var current = this.getStatus();
+            final var projected = this.changingStatusTo();
+            needConsumption = targetStatus.ordinal() <= (projected == null ? current.ordinal() : Math.min(current.ordinal(), projected.ordinal()));
         }
 
         if (needConsumption) {
@@ -208,16 +210,21 @@ BusyRefCounter busyRefCounter() {
         return this.busyRefCounter;
     }
 
-    public void submitUpgradeAction(CancellationSignaller signaller, ItemStatus<K, V, Ctx> status) {
+    public void finishAction() {
+        assertOpen();
+        Pair<Cancellable, ItemStatus<K, V, Ctx>> current = this.runningAction.getAndSet(null);
+        Assertions.assertTrue(current != null, "No action is present when trying to finish an action");
+    }
+
+    public void submitAction(Cancellable cancellation, ItemStatus<K, V, Ctx> status) {
         assertOpen();
-        final boolean success = this.runningUpgradeAction.compareAndSet(null, Pair.of(signaller, status));
+        final var success = this.runningAction.compareAndSet(null, Pair.of(cancellation, status));
         Assertions.assertTrue(success, "Only one action can happen at a time");
-        signaller.addListener(unused -> this.runningUpgradeAction.set(null));
     }
 
-    public void tryCancelUpgradeAction() {
+    public void tryCancelAction() {
         assertOpen();
-        final Pair<CancellationSignaller, ItemStatus<K, V, Ctx>> signaller = this.runningUpgradeAction.get();
+        final Pair<Cancellable, ItemStatus<K, V, Ctx>> signaller = this.runningAction.get();
         if (signaller != null) {
             signaller.left().cancel();
         }