Let the BatchLoaderScheduler process all CF completions as a batch

Author: fxbonnet

src/main/java/org/dataloader/DataLoaderHelper.java

@@ -28,6 +28,7 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
 import java.util.concurrent.CompletionStage;
+import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.atomic.AtomicReference;
 
 import static java.util.Collections.emptyList;
@@ -332,9 +333,8 @@ private CompletableFuture<List<V>> dispatchQueueBatch(List<K> keys, List<Object>
                         return CompletableFuture.completedFuture(values);
                     }
 
-                    List<K> clearCacheKeys = new ArrayList<>();
-                    var batchLoaderScheduler = loaderOptions.getBatchLoaderScheduler();
-                    CompletableFuture<Void>[] scheduledCompletions = new CompletableFuture[keys.size()];
+                    Collection<K> clearCacheKeys = new ConcurrentLinkedQueue<>();
+                    List<Runnable> completeValueRunnables = new ArrayList<>();
                     for (int idx = 0; idx < queuedFutures.size(); idx++) {
                         K key = keys.get(idx);
                         V value = values.get(idx);
@@ -360,14 +360,17 @@ private CompletableFuture<List<V>> dispatchQueueBatch(List<K> keys, List<Object>
                                 future.complete(value);
                             }
                         };
-                        if(batchLoaderScheduler != null) {
-                            scheduledCompletions[idx] = batchLoaderScheduler.scheduleCompletion(completeValueRunnable, key, value);
-                        } else {
-                            scheduledCompletions[idx] = CompletableFutureKit.run(completeValueRunnable);
-                        }
+                        completeValueRunnables.add(completeValueRunnable);
+                    }
+                    CompletableFuture<Void> result;
+                    var batchLoaderScheduler = loaderOptions.getBatchLoaderScheduler();
+                    if(batchLoaderScheduler != null) {
+                        result = batchLoaderScheduler.scheduleCompletions(completeValueRunnables, keys, values);
+                    } else {
+                        result = CompletableFuture.allOf(completeValueRunnables.stream().map(CompletableFutureKit::run).toArray(CompletableFuture[]::new));
                     }
                     // Wait for all completions to return
-                    return allOf(scheduledCompletions).thenApply(ignored -> {
+                    return result.thenApply(ignored -> {
                         possiblyClearCacheEntriesOnExceptions(clearCacheKeys);
                         return values;
                     });
@@ -392,7 +395,7 @@ private void assertResultSize(List<K> keys, List<V> values) {
         assertState(keys.size() == values.size(), () -> "The size of the promised values MUST be the same size as the key list");
     }
 
-    private void possiblyClearCacheEntriesOnExceptions(List<K> keys) {
+    private void possiblyClearCacheEntriesOnExceptions(Collection<K> keys) {
         if (keys.isEmpty()) {
             return;
         }

src/main/java/org/dataloader/scheduler/BatchLoaderScheduler.java

@@ -98,17 +98,17 @@ interface ScheduledBatchPublisherCall {
     /**
      * Schedules the completion of a {@link DataLoader} value's {@link java.util.concurrent.CompletableFuture} after a batch load.
      * <p>
-     * Override this to offload completions to another thread, avoiding slow chained work
-     * (e.g. {@code thenApply}) from executing inline on the dispatch thread.
+     * Override this to offload completions to another thread or execute in parallel using multiple threads, avoiding
+     * slow chained work (e.g. {@code thenApply}) from executing inline on the dispatch thread.
      * <p>
      * By default, completions run synchronously on the current thread, sequentially for each key in the batch.
      *
-     * @param completeValueRunnable the runnable that completes the value
-     * @param key                   the key being completed
-     * @param value                 the value returned by the batch function for this key
-     * @return a {@link CompletionStage} representing the scheduled work
+     * @param completeValueRunnables the runnables that completes each value
+     * @param keys                   the keys being completed
+     * @param values                 the values returned by the batch function for each key
+     * @return a {@link CompletionStage} representing the completion of all the scheduled work
      */
-    default <K, V> CompletableFuture<Void> scheduleCompletion(Runnable completeValueRunnable, K key, V value) {
-        return CompletableFutureKit.run(completeValueRunnable);
+    default <K, V> CompletableFuture<Void> scheduleCompletions(List<Runnable> completeValueRunnables, List<K> keys, List<V> values) {
+        return CompletableFuture.allOf(completeValueRunnables.stream().map(CompletableFutureKit::run).toArray(CompletableFuture[]::new));
     }
 }