Control parallelism on a separate virtual thread (#872)

pivovarit · May 1, 2024 · 5da805b · 5da805b
1 parent eb933ab
commit 5da805b
Show file tree

Hide file tree

Showing 4 changed files with 116 additions and 41 deletions.
diff --git a/src/main/java/com/pivovarit/collectors/AsyncParallelCollector.java b/src/main/java/com/pivovarit/collectors/AsyncParallelCollector.java
@@ -55,12 +55,21 @@ public BinaryOperator<List<CompletableFuture<R>>> combiner() {
 
     @Override
     public BiConsumer<List<CompletableFuture<R>>, T> accumulator() {
-        return (acc, e) -> acc.add(dispatcher.enqueue(() -> task.apply(e)));
+        return (acc, e) -> {
+            if (!dispatcher.isRunning()) {
+                dispatcher.start();
+            }
+            acc.add(dispatcher.enqueue(() -> task.apply(e)));
+        };
     }
 
     @Override
     public Function<List<CompletableFuture<R>>, CompletableFuture<C>> finisher() {
-        return futures -> combine(futures).thenApply(finalizer);
+        return futures -> {
+            dispatcher.stop();
+
+            return combine(futures).thenApply(finalizer);
+        };
     }
 
     @Override
@@ -103,7 +112,7 @@ private static <T> CompletableFuture<Stream<T>> combine(List<CompletableFuture<T
         requireNonNull(collector, "collector can't be null");
         requireNonNull(mapper, "mapper can't be null");
 
-        return new AsyncParallelCollector<>(mapper, Dispatcher.virtual(),s -> s.collect(collector));
+        return new AsyncParallelCollector<>(mapper, Dispatcher.virtual(), s -> s.collect(collector));
     }
 
     static <T, R, RR> Collector<T, ?, CompletableFuture<RR>> collectingWithCollector(Collector<R, ?, RR> collector, Function<T, R> mapper, Executor executor, int parallelism) {

diff --git a/src/main/java/com/pivovarit/collectors/Dispatcher.java b/src/main/java/com/pivovarit/collectors/Dispatcher.java
@@ -1,27 +1,46 @@
 package com.pivovarit.collectors;
 
+import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Executor;
+import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.FutureTask;
+import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.Semaphore;
-import java.util.function.BiConsumer;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.Function;
 import java.util.function.Supplier;
 
 /**
  * @author Grzegorz Piwowarek
  */
 final class Dispatcher<T> {
 
+    private static final Runnable POISON_PILL = () -> System.out.println("Why so serious?");
+
     private final CompletableFuture<Void> completionSignaller = new CompletableFuture<>();
+
+    private final BlockingQueue<Runnable> workingQueue = new LinkedBlockingQueue<>();
+
+    private final ExecutorService dispatcher = Executors.newVirtualThreadPerTaskExecutor();
     private final Executor executor;
     private final Semaphore limiter;
 
+    private final AtomicBoolean started = new AtomicBoolean(false);
+
+    private volatile boolean shortCircuited = false;
+
     private Dispatcher() {
-        this.executor = Executors.newVirtualThreadPerTaskExecutor();
+        this.executor = defaultExecutorService();
         this.limiter = null;
     }
 
+    private Dispatcher(int permits) {
+        this.executor = defaultExecutorService();
+        this.limiter = new Semaphore(permits);
+    }
+
     private Dispatcher(Executor executor, int permits) {
         this.executor = executor;
         this.limiter = new Semaphore(permits);
@@ -35,65 +54,107 @@ static <T> Dispatcher<T> virtual() {
         return new Dispatcher<>();
     }
 
-    CompletableFuture<T> enqueue(Supplier<T> supplier) {
-        InterruptibleCompletableFuture<T> future = new InterruptibleCompletableFuture<>();
-        completionSignaller.whenComplete(shortcircuit(future));
+    void start() {
+        if (!started.getAndSet(true)) {
+            dispatcher.execute(() -> {
+                try {
+                    while (true) {
+                        try {
+                            if (limiter != null) {
+                                limiter.acquire();
+                            }
+                        } catch (InterruptedException e) {
+                            handle(e);
+                        }
+                        Runnable task;
+                        if ((task = workingQueue.take()) != POISON_PILL) {
+                            executor.execute(() -> {
+                                try {
+                                    task.run();
+                                } finally {
+                                    if (limiter != null) {
+                                        limiter.release();
+                                    }
+                                }
+                            });
+                        } else {
+                            break;
+                        }
+                    }
+                } catch (Throwable e) {
+                    handle(e);
+                }
+            });
+        }
+    }
+
+    void stop() {
         try {
-            executor.execute(completionTask(supplier, future));
-        } catch (Throwable e) {
+            workingQueue.put(POISON_PILL);
+        } catch (InterruptedException e) {
             completionSignaller.completeExceptionally(e);
-            return CompletableFuture.failedFuture(e);
+        } finally {
+            dispatcher.shutdown();
         }
+    }
+
+    boolean isRunning() {
+        return started.get();
+    }
+
+    CompletableFuture<T> enqueue(Supplier<T> supplier) {
+        InterruptibleCompletableFuture<T> future = new InterruptibleCompletableFuture<>();
+        workingQueue.add(completionTask(supplier, future));
+        completionSignaller.exceptionally(shortcircuit(future));
         return future;
     }
 
-    private FutureTask<T> completionTask(Supplier<T> supplier, InterruptibleCompletableFuture<T> future) {
-        FutureTask<T> task = new FutureTask<>(() -> {
-            if (!completionSignaller.isCompletedExceptionally()) {
-                try {
-                    if (limiter == null) {
-                        future.complete(supplier.get());
-                    } else {
-                        try {
-                            limiter.acquire();
-                            future.complete(supplier.get());
-                        } finally {
-                            limiter.release();
-                        }
-                    }
-                } catch (Throwable e) {
-                    completionSignaller.completeExceptionally(e);
+    private FutureTask<Void> completionTask(Supplier<T> supplier, InterruptibleCompletableFuture<T> future) {
+        FutureTask<Void> task = new FutureTask<>(() -> {
+            try {
+                if (!shortCircuited) {
+                    future.complete(supplier.get());
                 }
+            } catch (Throwable e) {
+                handle(e);
             }
         }, null);
         future.completedBy(task);
         return task;
     }
 
-    private static <T> BiConsumer<T, Throwable> shortcircuit(InterruptibleCompletableFuture<?> future) {
-        return (__, throwable) -> {
-            if (throwable != null) {
-                future.completeExceptionally(throwable);
-                future.cancel(true);
-            }
+    private void handle(Throwable e) {
+        shortCircuited = true;
+        completionSignaller.completeExceptionally(e);
+        dispatcher.shutdownNow();
+    }
+
+    private static Function<Throwable, Void> shortcircuit(InterruptibleCompletableFuture<?> future) {
+        return throwable -> {
+            future.completeExceptionally(throwable);
+            future.cancel(true);
+            return null;
         };
     }
 
     static final class InterruptibleCompletableFuture<T> extends CompletableFuture<T> {
 
-        private volatile FutureTask<T> backingTask;
+        private volatile FutureTask<?> backingTask;
 
-        private void completedBy(FutureTask<T> task) {
+        private void completedBy(FutureTask<Void> task) {
             backingTask = task;
         }
 
         @Override
         public boolean cancel(boolean mayInterruptIfRunning) {
-            var task = backingTask;
-            if (task != null) {
-                task.cancel(mayInterruptIfRunning);
+            if (backingTask != null) {
+                backingTask.cancel(mayInterruptIfRunning);
             }
             return super.cancel(mayInterruptIfRunning);
         }
     }
+
+    private static ExecutorService defaultExecutorService() {
+        return Executors.newVirtualThreadPerTaskExecutor();
+    }
 }
diff --git a/src/main/java/com/pivovarit/collectors/ParallelStreamCollector.java b/src/main/java/com/pivovarit/collectors/ParallelStreamCollector.java
@@ -57,7 +57,10 @@ public Supplier<List<CompletableFuture<R>>> supplier() {
 
     @Override
     public BiConsumer<List<CompletableFuture<R>>, T> accumulator() {
-        return (acc, e) -> acc.add(dispatcher.enqueue(() -> function.apply(e)));
+        return (acc, e) -> {
+            dispatcher.start();
+            acc.add(dispatcher.enqueue(() -> function.apply(e)));
+        };
     }
 
     @Override
@@ -70,7 +73,10 @@ public BinaryOperator<List<CompletableFuture<R>>> combiner() {
 
     @Override
     public Function<List<CompletableFuture<R>>, Stream<R>> finisher() {
-        return completionStrategy;
+        return acc -> {
+            dispatcher.stop();
+            return completionStrategy.apply(acc);
+        };
     }
 
     @Override

diff --git a/src/test/java/com/pivovarit/collectors/FunctionalTest.java b/src/test/java/com/pivovarit/collectors/FunctionalTest.java
@@ -25,7 +25,6 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicLong;
-import java.util.concurrent.atomic.LongAdder;
 import java.util.function.Function;
 import java.util.stream.Collector;
 import java.util.stream.IntStream;