Simplify work-stealing Future implementation

Author: marcphilipp

junit-platform-engine/src/main/java/org/junit/platform/engine/support/hierarchical/BlockingAwareFuture.java

@@ -17,48 +17,40 @@
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
-import java.util.function.Supplier;
 
 import org.jspecify.annotations.Nullable;
 
-class BlockingAwareFuture<T extends @Nullable Object> extends DelegatingFuture<T> {
+abstract class BlockingAwareFuture<T extends @Nullable Object> extends DelegatingFuture<T> {
 
-	private final BlockHandler handler;
-
-	BlockingAwareFuture(Future<T> delegate, BlockHandler handler) {
+	BlockingAwareFuture(Future<T> delegate) {
 		super(delegate);
-		this.handler = handler;
 	}
 
 	@Override
 	public T get() throws InterruptedException, ExecutionException {
 		if (delegate.isDone()) {
 			return delegate.get();
 		}
-		return handle(delegate::get);
+		return handleSafely(delegate::get);
 	}
 
 	@Override
 	public T get(long timeout, TimeUnit unit) throws InterruptedException, ExecutionException, TimeoutException {
 		if (delegate.isDone()) {
 			return delegate.get();
 		}
-		return handle(() -> delegate.get(timeout, unit));
+		return handleSafely(() -> delegate.get(timeout, unit));
 	}
 
-	private T handle(Callable<T> callable) {
+	private T handleSafely(Callable<T> callable) {
 		try {
-			return handler.handle(delegate::isDone, callable);
+			return handle(callable);
 		}
 		catch (Exception e) {
 			throw throwAsUncheckedException(e);
 		}
 	}
 
-	interface BlockHandler {
-
-		<T extends @Nullable Object> T handle(Supplier<Boolean> blockingUnnecessary, Callable<T> callable)
-				throws Exception;
+	protected abstract T handle(Callable<T> callable) throws Exception;
 
-	}
 }

junit-platform-engine/src/main/java/org/junit/platform/engine/support/hierarchical/ConcurrentHierarchicalTestExecutorService.java

@@ -96,8 +96,7 @@ public void close() {
 			return completedFuture(null);
 		}
 
-		var entry = enqueue(testTask);
-		return new BlockingAwareFuture<@Nullable Void>(entry.future(), new WorkerThread.BlockHandler(entry));
+		return new WorkStealingFuture(enqueue(testTask));
 	}
 
 	@Override
@@ -434,28 +433,36 @@ private interface BlockingAction<T> {
 			T run() throws InterruptedException;
 		}
 
-		private record BlockHandler(WorkQueue.Entry entry) implements BlockingAwareFuture.BlockHandler {
+	}
 
-			@Override
-			public <T> T handle(Supplier<Boolean> blockingUnnecessary, Callable<T> callable) throws Exception {
-				var workerThread = get();
-				if (workerThread == null || entry.future.isDone()) {
+	private static class WorkStealingFuture extends BlockingAwareFuture<@Nullable Void> {
+
+		private final WorkQueue.Entry entry;
+
+		WorkStealingFuture(WorkQueue.Entry entry) {
+			super(entry.future);
+			this.entry = entry;
+		}
+
+		@Override
+		protected @Nullable Void handle(Callable<@Nullable Void> callable) throws Exception {
+			var workerThread = WorkerThread.get();
+			if (workerThread == null || entry.future.isDone()) {
+				return callable.call();
+			}
+			workerThread.tryToStealWork(entry);
+			if (entry.future.isDone()) {
+				return callable.call();
+			}
+			LOGGER.trace(() -> "blocking for child task");
+			return workerThread.runBlocking(() -> {
+				try {
 					return callable.call();
 				}
-				workerThread.tryToStealWork(entry);
-				if (entry.future.isDone()) {
-					return callable.call();
+				catch (Exception ex) {
+					throw throwAsUncheckedException(ex);
 				}
-				LOGGER.trace(() -> "blocking for child task");
-				return workerThread.runBlocking(() -> {
-					try {
-						return callable.call();
-					}
-					catch (Exception ex) {
-						throw throwAsUncheckedException(ex);
-					}
-				});
-			}
+			});
 		}
 
 	}
@@ -612,5 +619,4 @@ void reacquire() throws InterruptedException {
 			reacquisitionToken = null;
 		}
 	}
-
 }