Refactor task callback handling for cancellation/continuation (#111)

* step 1

* step 2

Author: Simn

src/hxcoro/concurrent/ThreadSafeCallbacks.hx

@@ -0,0 +1,88 @@
+package hxcoro.concurrent;
+
+import hxcoro.concurrent.AtomicState;
+import hxcoro.concurrent.BackOff;
+
+enum abstract CallbacksState(Int) to Int {
+	final Ready;
+	final Modifying;
+	final Finished;
+}
+
+abstract class ThreadSafeCallbacks<Element, HandleIn : HandleOut, HandleOut> {
+	var handles:Null<Array<HandleIn>>;
+	final execute:HandleIn -> Void;
+	final state:AtomicState<CallbacksState>;
+
+	function new(execute:HandleIn -> Void) {
+		handles = null;
+		this.execute = execute;
+		state = new AtomicState(Ready);
+	}
+
+	// single threaded
+
+	public function run() {
+		while (true) {
+			switch (state.compareExchange(Ready, Finished)) {
+				case Ready:
+					break;
+				case Modifying:
+					BackOff.backOff();
+				case Finished:
+					// already done
+					return;
+			}
+		}
+		final handles = handles;
+		if (handles == null) {
+			return;
+		}
+		this.handles = null;
+		for (handle in handles) {
+			// TODO: should we catch errors from the callbacks here?
+			execute(handle);
+		}
+	}
+
+	// thread-safe
+
+	abstract function createHandle(element:Element):HandleIn;
+
+	public function add(element:Element):Null<HandleOut> {
+		final handle = createHandle(element);
+		while (true) {
+			switch (state.compareExchange(Ready, Modifying)) {
+				case Ready:
+					break;
+				case Modifying:
+					BackOff.backOff();
+				case Finished:
+					execute(handle);
+					return null;
+			}
+		}
+		handles ??= [];
+		handles.push(handle);
+		state.store(Ready);
+		return handle;
+	}
+
+	public function remove(handle:HandleIn) {
+		while (true) {
+			switch (state.compareExchange(Ready, Modifying)) {
+				case Ready:
+					break;
+				case Modifying:
+					BackOff.backOff();
+				case Finished:
+					// already cleared, nothing to do
+					return;
+			}
+		}
+		if (handles != null) {
+			handles.remove(handle);
+		}
+		state.store(Ready);
+	}
+}

src/hxcoro/task/AbstractTask.hx

@@ -1,13 +1,13 @@
 package hxcoro.task;
 
-import hxcoro.concurrent.AtomicObject;
-import hxcoro.concurrent.AtomicState;
-import hxcoro.concurrent.AtomicInt;
+import haxe.Exception;
 import haxe.coro.cancellation.ICancellationCallback;
 import haxe.coro.cancellation.ICancellationHandle;
 import haxe.coro.cancellation.ICancellationToken;
 import haxe.exceptions.CancellationException;
-import haxe.Exception;
+import hxcoro.concurrent.AtomicInt;
+import hxcoro.concurrent.AtomicObject;
+import hxcoro.concurrent.AtomicState;
 
 @:using(AbstractTask.TaskStateTools)
 enum abstract TaskState(Int) to Int {
@@ -155,7 +155,7 @@ abstract class AbstractTask implements ICancellationToken {
 	}
 
 	public function onCancellationRequested(callback:ICancellationCallback):ICancellationHandle {
-		return cancellationManager.addCallback(callback);
+		return cancellationManager.add(callback);
 	}
 
 	/**

src/hxcoro/task/CoroBaseTask.hx

@@ -18,6 +18,7 @@ import haxe.coro.context.Key;
 import haxe.coro.context.IElement;
 import haxe.coro.dispatchers.Dispatcher;
 import haxe.coro.cancellation.CancellationToken;
+import hxcoro.concurrent.ThreadSafeCallbacks;
 
 private class CoroTaskWith<T> implements ICoroNodeWith {
 	public var context(get, null):Context;
@@ -51,64 +52,13 @@ private class CoroTaskWith<T> implements ICoroNodeWith {
 		return task.without(...keys);
 	}
 }
-enum abstract AggregatorState(Int) to Int {
-	final Ready;
-	final Modifying;
-	final Finished;
-}
-
-class ThreadSafeAggregator<T> {
-	var array:Null<Array<T>>;
-	final func:T -> Void;
-	final state:AtomicState<AggregatorState>;
-
-	public function new(func:T -> Void) {
-		array = null;
-		this.func = func;
-		state = new AtomicState(Ready);
-	}
-
-	// single threaded
-
-	public function run() {
-		while (true) {
-			switch (state.compareExchange(Ready, Finished)) {
-				case Ready:
-					break;
-				case Modifying:
-					BackOff.backOff();
-				case Finished:
-					// already done
-					return;
-			}
-		}
-		final array = array;
-		if (array == null) {
-			return;
-		}
-		this.array = null;
-		for (element in array) {
-			func(element);
-		}
+class TaskContinuationManager extends ThreadSafeCallbacks<IContinuation<Any>, IContinuation<Any>, IContinuation<Any>> {
+	public function new(task:CoroBaseTask<Any>) {
+		super(handle -> handle.resume(task.get(), task.getError()));
 	}
 
-	// thread-safe
-
-	public function add(element:T) {
-		while (true) {
-			switch (state.compareExchange(Ready, Modifying)) {
-				case Ready:
-					break;
-				case Modifying:
-					BackOff.backOff();
-				case Finished:
-					func(element);
-					return;
-			}
-		}
-		array ??= [];
-		array.push(element);
-		state.store(Ready);
+	function createHandle(element:IContinuation<Any>) {
+		return element;
 	}
 }
 
@@ -124,7 +74,7 @@ abstract class CoroBaseTask<T> extends AbstractTask implements ICoroNode impleme
 	public var context(get, null):Context;
 
 	final nodeStrategy:INodeStrategy;
-	final awaitingContinuations:ThreadSafeAggregator<IContinuation<T>>;
+	final awaitingContinuations:TaskContinuationManager;
 	final awaitingChildContinuation:AtomicObject<Null<IContinuation<Any>>>;
 	var result:Null<T>;
 
@@ -135,9 +85,7 @@ abstract class CoroBaseTask<T> extends AbstractTask implements ICoroNode impleme
 		final parent = context.get(CoroBaseTask);
 		this.context = context.clone().with(this).set(CancellationToken, this);
 		this.nodeStrategy = nodeStrategy;
-		awaitingContinuations = new ThreadSafeAggregator<IContinuation<T>>(cont ->
-			cont.resume(result, error.load())
-		);
+		awaitingContinuations = new TaskContinuationManager(this);
 		awaitingChildContinuation = new AtomicObject(null);
 		super(parent, initialState);
 	}

src/hxcoro/task/TaskCancellationManager.hx

@@ -1,25 +1,16 @@
 package hxcoro.task;
 
-import hxcoro.concurrent.AtomicState;
-import hxcoro.concurrent.BackOff;
 import haxe.coro.cancellation.ICancellationCallback;
 import haxe.coro.cancellation.ICancellationHandle;
-
-private class NoOpCancellationHandle implements ICancellationHandle {
-	public function new() {}
-
-	public function close() {}
-}
+import hxcoro.concurrent.AtomicState;
+import hxcoro.concurrent.ThreadSafeCallbacks;
 
 private enum abstract HandleState(Int) to Int {
 	final Open;
 	final Closed;
 }
 
 class CancellationHandle implements ICancellationHandle {
-
-	static public final noOpCancellationHandle = new NoOpCancellationHandle();
-
 	final callback:ICancellationCallback;
 	final manager:TaskCancellationManager;
 	var closed:AtomicState<HandleState>;
@@ -30,12 +21,12 @@ class CancellationHandle implements ICancellationHandle {
 		closed = new AtomicState(Open);
 	}
 
-	public function run() {
+	public function run(task:AbstractTask) {
 		if (closed.compareExchange(Open, Closed) != Open) {
 			return;
 		}
 
-		final error = manager.task.getError();
+		final error = task.getError();
 		callback.onCancellation(error.orCancellationException());
 	}
 
@@ -46,85 +37,12 @@ class CancellationHandle implements ICancellationHandle {
 	}
 }
 
-private enum abstract ManagerState(Int) to Int {
-	final Ready;
-	final Modifying;
-	final Finished;
-}
-
-class TaskCancellationManager {
-	public final task:AbstractTask;
-	final state:AtomicState<ManagerState>;
-	var handles:Null<Array<CancellationHandle>>;
-
+class TaskCancellationManager extends ThreadSafeCallbacks<ICancellationCallback, CancellationHandle, ICancellationHandle> {
 	public function new(task:AbstractTask) {
-		this.task = task;
-		handles = null;
-		state = new AtomicState(Ready);
-	}
-
-	// single-threaded
-
-	public function run() {
-		while (true) {
-			switch (state.compareExchange(Ready, Finished)) {
-				case Ready:
-					break;
-				case Modifying:
-					BackOff.backOff();
-				case Finished:
-					// already done
-					return;
-			}
-		}
-		final handles = handles;
-		if (handles == null) {
-			return;
-		}
-		this.handles = null;
-		for (handle in handles) {
-			// TODO: should we catch errors from the callbacks here?
-			handle.run();
-		}
-
+		super(handle -> handle.run(task));
 	}
 
-	// thread-safe
-
-	public function addCallback(callback:ICancellationCallback):ICancellationHandle {
-		final handle = new CancellationHandle(callback, this);
-		while (true) {
-			switch (state.compareExchange(Ready, Modifying)) {
-				case Ready:
-					break;
-				case Modifying:
-					BackOff.backOff();
-				case Finished:
-					handle.run();
-					return CancellationHandle.noOpCancellationHandle;
-			}
-		}
-		handles ??= [];
-		handles.push(handle);
-		state.store(Ready);
-		return handle;
-	}
-
-	public function remove(handle:CancellationHandle) {
-		while (true) {
-			switch (state.compareExchange(Ready, Modifying)) {
-				case Ready:
-					break;
-				case Modifying:
-					BackOff.backOff();
-				case Finished:
-					// already cleared, nothing to do
-					return;
-			}
-		}
-		if (handles != null) {
-			handles.remove(handle);
-		}
-		state.store(Ready);
+	function createHandle(element:ICancellationCallback) {
+		return new CancellationHandle(element, this);
 	}
 }