From 848fa884094a488ba304f4c583f94f26be92b2e6 Mon Sep 17 00:00:00 2001
From: Grzegorz Piwowarek <gpiwowarek@gmail.com>
Date: Sun, 28 Jan 2024 15:28:53 +0100
Subject: [PATCH] Seal CompletionStrategy

---
 .../collectors/AsyncParallelCollector.java    |  1 -
 .../collectors/CompletionStrategy.java        | 23 ++++++++++++++++---
 .../com/pivovarit/collectors/Dispatcher.java  | 16 +++++--------
 .../collectors/FutureCollectors.java          |  6 ++---
 .../collectors/ArchitectureTest.java          |  1 +
 5 files changed, 29 insertions(+), 18 deletions(-)

diff --git a/src/main/java/com/pivovarit/collectors/AsyncParallelCollector.java b/src/main/java/com/pivovarit/collectors/AsyncParallelCollector.java
index bd6a782b..7af367c9 100644
--- a/src/main/java/com/pivovarit/collectors/AsyncParallelCollector.java
+++ b/src/main/java/com/pivovarit/collectors/AsyncParallelCollector.java
@@ -7,7 +7,6 @@
 import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Executor;
-import java.util.concurrent.Executors;
 import java.util.function.BiConsumer;
 import java.util.function.BinaryOperator;
 import java.util.function.Function;
diff --git a/src/main/java/com/pivovarit/collectors/CompletionStrategy.java b/src/main/java/com/pivovarit/collectors/CompletionStrategy.java
index bd74edcd..9feb90f1 100644
--- a/src/main/java/com/pivovarit/collectors/CompletionStrategy.java
+++ b/src/main/java/com/pivovarit/collectors/CompletionStrategy.java
@@ -6,13 +6,30 @@
 import java.util.stream.Stream;
 import java.util.stream.StreamSupport;
 
-interface CompletionStrategy<T> extends Function<List<CompletableFuture<T>>, Stream<T>> {
+sealed interface CompletionStrategy<T> extends Function<List<CompletableFuture<T>>, Stream<T>> permits CompletionStrategy.Unordered, CompletionStrategy.Ordered {
+
+    Unordered<?> UNORDERED = new Unordered<>();
+    Ordered<?> ORDERED = new Ordered<>();
 
     static <R> CompletionStrategy<R> unordered() {
-        return futures -> StreamSupport.stream(new CompletionOrderSpliterator<>(futures), false);
+        return (CompletionStrategy<R>) UNORDERED;
     }
 
     static <R> CompletionStrategy<R> ordered() {
-        return futures -> futures.stream().map(CompletableFuture::join);
+        return (CompletionStrategy<R>) ORDERED;
+    }
+
+    final class Unordered<T> implements CompletionStrategy<T> {
+        @Override
+        public Stream<T> apply(List<CompletableFuture<T>> futures) {
+            return StreamSupport.stream(new CompletionOrderSpliterator<>(futures), false);
+        }
+    }
+
+    final class Ordered<T> implements CompletionStrategy<T> {
+        @Override
+        public Stream<T> apply(List<CompletableFuture<T>> futures) {
+            return futures.stream().map(CompletableFuture::join);
+        }
     }
 }
diff --git a/src/main/java/com/pivovarit/collectors/Dispatcher.java b/src/main/java/com/pivovarit/collectors/Dispatcher.java
index f9d1f0d1..dffa9a8d 100644
--- a/src/main/java/com/pivovarit/collectors/Dispatcher.java
+++ b/src/main/java/com/pivovarit/collectors/Dispatcher.java
@@ -54,7 +54,12 @@ private FutureTask<T> completionTask(Supplier<T> supplier, InterruptibleCompleta
                     if (limiter == null) {
                         future.complete(supplier.get());
                     } else {
-                        withLimiter(supplier, future);
+                        try {
+                            limiter.acquire();
+                            future.complete(supplier.get());
+                        } finally {
+                            limiter.release();
+                        }
                     }
                 } catch (Throwable e) {
                     completionSignaller.completeExceptionally(e);
@@ -65,15 +70,6 @@ private FutureTask<T> completionTask(Supplier<T> supplier, InterruptibleCompleta
         return task;
     }
 
-    private void withLimiter(Supplier<T> supplier, InterruptibleCompletableFuture<T> future) throws InterruptedException {
-        try {
-            limiter.acquire();
-            future.complete(supplier.get());
-        } finally {
-            limiter.release();
-        }
-    }
-
     private static <T> BiConsumer<T, Throwable> shortcircuit(InterruptibleCompletableFuture<?> future) {
         return (__, throwable) -> {
             if (throwable != null) {
diff --git a/src/main/java/com/pivovarit/collectors/FutureCollectors.java b/src/main/java/com/pivovarit/collectors/FutureCollectors.java
index 16156ac9..afab87b4 100644
--- a/src/main/java/com/pivovarit/collectors/FutureCollectors.java
+++ b/src/main/java/com/pivovarit/collectors/FutureCollectors.java
@@ -13,14 +13,12 @@
 final class FutureCollectors {
     static <T, R> Collector<CompletableFuture<T>, ?, CompletableFuture<R>> toFuture(Collector<T, ?, R> collector) {
         return Collectors.collectingAndThen(toList(), list -> {
-            CompletableFuture<R> future = CompletableFuture
-              .allOf(list.toArray(new CompletableFuture[0]))
+            var future = CompletableFuture.allOf(list.toArray(CompletableFuture[]::new))
               .thenApply(__ -> list.stream()
                 .map(CompletableFuture::join)
                 .collect(collector));
 
-            // CompletableFuture#allOf doesn't shortcircuit on exception so that requires manual handling
-            for (CompletableFuture<T> f : list) {
+            for (var f : list) {
                 f.whenComplete((t, throwable) -> {
                     if (throwable != null) {
                         future.completeExceptionally(throwable);
diff --git a/src/test/java/com/pivovarit/collectors/ArchitectureTest.java b/src/test/java/com/pivovarit/collectors/ArchitectureTest.java
index 75e06af3..29dc68db 100644
--- a/src/test/java/com/pivovarit/collectors/ArchitectureTest.java
+++ b/src/test/java/com/pivovarit/collectors/ArchitectureTest.java
@@ -16,6 +16,7 @@ class ArchitectureTest {
     @ArchTest
     static final ArchRule shouldHaveSingleFacade = classes()
       .that().arePublic()
+      .and().areNotNestedClasses()
       .should().haveSimpleName("ParallelCollectors").orShould().haveSimpleName("Batching")
       .andShould().haveOnlyFinalFields()
       .andShould().haveOnlyPrivateConstructors()