Dispatcher to use caller thread instead of dedicated scheduler thread

pivovarit · Oct 10, 2023 · fddaf74 · fddaf74
1 parent 15126c4
commit fddaf74
Show file tree

Hide file tree

Showing 5 changed files with 33 additions and 127 deletions.
diff --git a/src/main/java/com/pivovarit/collectors/AsyncParallelCollector.java b/src/main/java/com/pivovarit/collectors/AsyncParallelCollector.java
@@ -55,21 +55,12 @@ public BinaryOperator<List<CompletableFuture<R>>> combiner() {
 
     @Override
     public BiConsumer<List<CompletableFuture<R>>, T> accumulator() {
-        return (acc, e) -> {
-            if (!dispatcher.isRunning()) {
-                dispatcher.start();
-            }
-            acc.add(dispatcher.enqueue(() -> mapper.apply(e)));
-        };
+        return (acc, e) -> acc.add(dispatcher.enqueue(() -> mapper.apply(e)));
     }
 
     @Override
     public Function<List<CompletableFuture<R>>, CompletableFuture<C>> finisher() {
-        return futures -> {
-            dispatcher.stop();
-
-            return combine(futures).thenApply(processor);
-        };
+        return futures -> combine(futures).thenApply(processor);
     }
 
     @Override

diff --git a/src/main/java/com/pivovarit/collectors/Dispatcher.java b/src/main/java/com/pivovarit/collectors/Dispatcher.java
@@ -1,37 +1,23 @@
 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.concurrent.SynchronousQueue;
-import java.util.concurrent.ThreadPoolExecutor;
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicBoolean;
-import java.util.function.Function;
+import java.util.function.BiConsumer;
 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 = newLazySingleThreadExecutor();
     private final Executor executor;
     private final Semaphore limiter;
 
-    private final AtomicBoolean started = new AtomicBoolean(false);
-
     private volatile boolean shortCircuited = false;
 
     private Dispatcher(int permits) {
@@ -52,61 +38,20 @@ static <T> Dispatcher<T> virtual(int permits) {
         return new Dispatcher<>(permits);
     }
 
-    void start() {
-        if (!started.getAndSet(true)) {
-            dispatcher.execute(() -> {
-                try {
-                    while (true) {
-                        Runnable task;
-                        if ((task = workingQueue.take()) != POISON_PILL) {
-                            executor.execute(() -> {
-                                try {
-                                    limiter.acquire();
-                                    task.run();
-                                } catch (InterruptedException e) {
-                                    handle(e);
-                                } finally {
-                                    limiter.release();
-                                }
-                            });
-                        } else {
-                            break;
-                        }
-                    }
-                } catch (Throwable e) {
-                    handle(e);
-                }
-            });
-        }
-    }
-
-    void stop() {
-        try {
-            workingQueue.put(POISON_PILL);
-        } catch (InterruptedException e) {
-            completionSignaller.completeExceptionally(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));
+        executor.execute(completionTask(supplier, future));
+        completionSignaller.whenComplete(shortcircuit(future));
         return future;
     }
 
     private FutureTask<Void> completionTask(Supplier<T> supplier, InterruptibleCompletableFuture<T> future) {
         FutureTask<Void> task = new FutureTask<>(() -> {
+            if (shortCircuited) {
+                return;
+            }
             try {
-                if (!shortCircuited) {
-                    future.complete(supplier.get());
-                }
+                withLimiter(supplier, future);
             } catch (Throwable e) {
                 handle(e);
             }
@@ -115,46 +60,48 @@ private FutureTask<Void> completionTask(Supplier<T> supplier, InterruptibleCompl
         return task;
     }
 
+
+    private void withLimiter(Supplier<T> supplier, InterruptibleCompletableFuture<T> future) throws InterruptedException {
+        try {
+            limiter.acquire();
+            future.complete(supplier.get());
+        } finally {
+            limiter.release();
+        }
+    }
+
     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;
+    private static <T> BiConsumer<T, Throwable> shortcircuit(InterruptibleCompletableFuture<?> future) {
+        return (__, throwable) -> {
+            if (throwable != null) {
+                future.completeExceptionally(throwable);
+                future.cancel(true);
+            }
         };
     }
 
-    private static ThreadPoolExecutor newLazySingleThreadExecutor() {
-        return new ThreadPoolExecutor(1, 1,
-          0L, TimeUnit.MILLISECONDS,
-          new SynchronousQueue<>(), // dispatcher always executes a single task
-          Thread.ofPlatform()
-            .name("parallel-collectors-dispatcher-", 0)
-            .daemon(false)
-            .factory());
-    }
-
     static final class InterruptibleCompletableFuture<T> extends CompletableFuture<T> {
 
         private volatile FutureTask<?> backingTask;
+
         private void completedBy(FutureTask<Void> task) {
             backingTask = task;
         }
 
         @Override
         public boolean cancel(boolean mayInterruptIfRunning) {
-            if (backingTask != null) {
-                backingTask.cancel(mayInterruptIfRunning);
+            var task = backingTask;
+            if (task != null) {
+                task.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,10 +57,7 @@ public Supplier<List<CompletableFuture<R>>> supplier() {
 
     @Override
     public BiConsumer<List<CompletableFuture<R>>, T> accumulator() {
-        return (acc, e) -> {
-            dispatcher.start();
-            acc.add(dispatcher.enqueue(() -> function.apply(e)));
-        };
+        return (acc, e) -> acc.add(dispatcher.enqueue(() -> function.apply(e)));
     }
 
     @Override
@@ -73,10 +70,7 @@ public BinaryOperator<List<CompletableFuture<R>>> combiner() {
 
     @Override
     public Function<List<CompletableFuture<R>>, Stream<R>> finisher() {
-        return acc -> {
-            dispatcher.stop();
-            return completionStrategy.apply(acc);
-        };
+        return completionStrategy;
     }
 
     @Override

diff --git a/src/test/java/com/pivovarit/collectors/CompletionOrderSpliteratorTest.java b/src/test/java/com/pivovarit/collectors/CompletionOrderSpliteratorTest.java
@@ -103,7 +103,7 @@ void shouldRestoreInterrupt() throws InterruptedException {
                 spliterator.tryAdvance(i -> {});
             } catch (Exception e) {
                 while (true) {
-
+                    Thread.onSpinWait();
                 }
             }
         });

diff --git a/src/test/java/com/pivovarit/collectors/FunctionalTest.java b/src/test/java/com/pivovarit/collectors/FunctionalTest.java
@@ -162,7 +162,6 @@ private static <R extends Collection<Integer>> Stream<DynamicTest> tests(Collect
           shouldCollectNElementsWithNParallelism(collector, name, PARALLELISM),
           shouldCollectToEmpty(collector, name),
           shouldStartConsumingImmediately(collector, name),
-          shouldTerminateAfterConsumingAllElements(collector, name),
           shouldNotBlockTheCallingThread(collector, name),
           shouldRespectParallelism(collector, name),
           shouldHandleThrowable(collector, name),
@@ -184,7 +183,6 @@ private static <R extends Collection<Integer>> Stream<DynamicTest> streamingTest
           shouldCollect(collector, name, PARALLELISM),
           shouldCollectToEmpty(collector, name),
           shouldStartConsumingImmediately(collector, name),
-          shouldTerminateAfterConsumingAllElements(collector, name),
           shouldNotBlockTheCallingThread(collector, name),
           shouldRespectParallelism(collector, name),
           shouldHandleThrowable(collector, name),
@@ -286,30 +284,6 @@ private static <R extends Collection<Integer>> DynamicTest shouldCollectNElement
         });
     }
 
-    private static <R extends Collection<Integer>> DynamicTest shouldTerminateAfterConsumingAllElements(CollectorSupplier<Function<Integer, Integer>, Executor, Integer, Collector<Integer, ?, CompletableFuture<R>>> factory, String name) {
-        return dynamicTest(format("%s: should terminate after consuming all elements", name), () -> {
-            List<Integer> elements = IntStream.range(0, 10).boxed().collect(toList());
-            Collector<Integer, ?, CompletableFuture<R>> ctor = factory.apply(i -> i, executor, 10);
-            Collection<Integer> result = elements.stream().collect(ctor)
-              .join();
-
-            assertThat(result).hasSameElementsAs(elements);
-
-            if (ctor instanceof AsyncParallelCollector) {
-                Field dispatcherField = AsyncParallelCollector.class.getDeclaredField("dispatcher");
-                dispatcherField.setAccessible(true);
-                Dispatcher<?> dispatcher = (Dispatcher<?>) dispatcherField.get(ctor);
-                Field innerDispatcherField = Dispatcher.class.getDeclaredField("dispatcher");
-                innerDispatcherField.setAccessible(true);
-                ExecutorService executor = (ExecutorService) innerDispatcherField.get(dispatcher);
-
-                await()
-                  .atMost(Duration.ofSeconds(2))
-                  .until(executor::isTerminated);
-            }
-        });
-    }
-
     private static <R extends Collection<Integer>> DynamicTest shouldMaintainOrder(CollectorSupplier<Function<Integer, Integer>, Executor, Integer, Collector<Integer, ?, CompletableFuture<R>>> collector, String name) {
         return dynamicTest(format("%s: should maintain order", name), () -> {
             int parallelism = 4;