Allow BatchLoaderScheduler process all CF completions as a batch and potentially in parallel

Author: fxbonnet

README.md

@@ -655,15 +655,17 @@ Perhaps you want this completion to happen more asynchronously so that the `.dis
 quickly and is not bound to the `.doSomethingSlow(v)` call.
 
 By default, the dispatch completion is done on the current thread in a synchronous manner, which will include
-any extra `CompletableFuture` chained methods.
+any extra `CompletableFuture` chained methods. If multiple keys were loaded, the dispatch completion is done
+sequentially.
 
 This is an example of running the completion step in an asynchronous manner :
 
 ```java
 
 @Override
-public <K> CompletionStage<Void> scheduleCompletion(Runnable completeValuesRunnable, List<K> keys, BatchLoaderEnvironment environment) {
-    return CompletableFuture.runAsync(completeValuesRunnable);
+public <K> CompletionStage<Void> scheduleCompletion(List<Runnable> completeValueRunnables, List<K> keys, BatchLoaderEnvironment environment) {
+    return CompletableFuture.allOf(completeValueRunnables.stream()
+            .map(CompletableFuture::runAsync).toArray(CompletableFuture[]::new));
 }
 ```
 

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

@@ -19,6 +19,7 @@
 import java.time.Instant;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collection;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
@@ -27,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;
@@ -329,17 +331,18 @@ private CompletableFuture<List<V>> dispatchQueueBatch(List<K> keys, List<Object>
                         return CompletableFutureKit.success(values);
                     }
 
-                    Runnable completeValuesRunnable = () -> {
-                        List<K> clearCacheKeys = new ArrayList<>();
-                        for (int idx = 0; idx < queuedFutures.size(); idx++) {
-                            K key = keys.get(idx);
-                            V value = values.get(idx);
-                            Object callContext = keyContexts.get(idx);
-                            CompletableFuture<V> future = queuedFutures.get(idx);
+                    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);
+                        Object callContext = keyContexts.get(idx);
+                        CompletableFuture<V> future = queuedFutures.get(idx);
+                        Runnable completeValueRunnable = () -> {
                             if (value instanceof Throwable) {
                                 stats.incrementLoadErrorCount(new IncrementLoadErrorCountStatisticsContext<>(key, callContext));
                                 future.completeExceptionally((Throwable) value);
-                                clearCacheKeys.add(keys.get(idx));
+                                clearCacheKeys.add(key);
                             } else if (value instanceof Try) {
                                 // we allow the batch loader to return a Try so we can better represent a computation
                                 // that might have worked or not.
@@ -349,16 +352,18 @@ private CompletableFuture<List<V>> dispatchQueueBatch(List<K> keys, List<Object>
                                 } else {
                                     stats.incrementLoadErrorCount(new IncrementLoadErrorCountStatisticsContext<>(key, callContext));
                                     future.completeExceptionally(tryValue.getThrowable());
-                                    clearCacheKeys.add(keys.get(idx));
+                                    clearCacheKeys.add(key);
                                 }
                             } else {
                                 future.complete(value);
                             }
-                        }
+                        };
+                        completeValueRunnables.add(completeValueRunnable);
+                    }
+                    return scheduleCompletion(environment, keys, values, completeValueRunnables).thenApply(ignored -> {
                         possiblyClearCacheEntriesOnExceptions(clearCacheKeys);
-                    };
-
-                    return scheduleCompletion(environment, keys, values, completeValuesRunnable);
+                        return values;
+                    });
                 }).exceptionally(ex -> {
                     stats.incrementBatchLoadExceptionCount(new IncrementBatchLoadExceptionCountStatisticsContext<>(keys, keyContexts));
                     if (ex instanceof CompletionException) {
@@ -375,14 +380,14 @@ private CompletableFuture<List<V>> dispatchQueueBatch(List<K> keys, List<Object>
                 });
     }
 
-    private CompletableFuture<List<V>> scheduleCompletion(BatchLoaderEnvironment environment, List<K> keys, List<V> values, Runnable completeValuesRunnable) {
+    private CompletableFuture<List<V>> scheduleCompletion(BatchLoaderEnvironment environment, List<K> keys, List<V> values, List<Runnable> completeValueRunnables) {
         BatchLoaderScheduler batchLoaderScheduler = loaderOptions.getBatchLoaderScheduler();
         CompletionStage<?> scheduledCompletion;
         if (batchLoaderScheduler != null) {
             scheduledCompletion = batchLoaderScheduler
-                    .scheduleCompletion(completeValuesRunnable, keys, environment);
+                    .scheduleCompletion(completeValueRunnables, keys, environment);
         } else {
-            scheduledCompletion = CompletableFutureKit.run(completeValuesRunnable);
+            scheduledCompletion = CompletableFutureKit.runAll(completeValueRunnables);
         }
         return scheduledCompletion
                 .thenApply(ignored -> values)
@@ -400,7 +405,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/impl/CompletableFutureKit.java

@@ -84,6 +84,13 @@ public static <T> CompletableFuture<List<T>> allOfFlatMap(List<CompletableFuture
                         .collect(toList()));
     }
 
+    /**
+     * Runs the given {@link Runnable} synchronously on the current thread, returning a
+     * {@link CompletableFuture} that is completed normally or exceptionally based on the outcome.
+     *
+     * @param runnable the task to execute
+     * @return a completed future, or a failed future if the runnable throws
+     */
     public static CompletableFuture<Void> run(Runnable runnable) {
         try {
             runnable.run();
@@ -92,4 +99,16 @@ public static CompletableFuture<Void> run(Runnable runnable) {
             return CompletableFutureKit.failedFuture(e);
         }
     }
+
+    /**
+     * Runs all the {@link Runnable} from the list synchronously on the current thread, returning a
+     * {@link CompletableFuture} that is completed normally or exceptionally based on the outcome.
+     *
+     * @param runnables the list of tasks to execute
+     * @return a completed future, or a failed future if any of the tasks throws
+     */
+    public static CompletableFuture<Void> runAll(List<Runnable> runnables) {
+        return CompletableFuture.allOf(runnables.stream().map(CompletableFutureKit::run).toArray(CompletableFuture[]::new));
+    }
+
 }

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

@@ -121,15 +121,15 @@ interface ScheduledBatchPublisherCall {
      * By default, the dispatch completion is done on the current thread in a synchronous manner, which will include
      * any extra {@link java.util.concurrent.CompletableFuture} dependent chained methods.
      *
-     * @param completeValuesRunnable this is the runnable that the {@link DataLoader} engine code needs to be run
+     * @param completeValueRunnables these are the runnable tasks that the {@link DataLoader} engine code needs to be run
      * @param keys          this is the list of keys that will be passed to the {@link BatchPublisher}.
      *                      This is provided only for informative reasons and, you can't change the keys that are used
      * @param environment   this is the {@link BatchLoaderEnvironment} in place,
      *
      * @return a {@link CompletionStage} representing this work is being scheduled
      */
-    default <K>  CompletionStage<Void> scheduleCompletion(Runnable completeValuesRunnable, List<K> keys, BatchLoaderEnvironment environment) {
-        return CompletableFutureKit.run(completeValuesRunnable);
+    default <K> CompletionStage<Void> scheduleCompletion(List<Runnable> completeValueRunnables, List<K> keys, BatchLoaderEnvironment environment) {
+        return CompletableFutureKit.runAll(completeValueRunnables);
     }
 
 }

src/test/java/org/dataloader/impl/CompletableFutureKitTest.java

@@ -99,4 +99,31 @@ void run() {
 
 
     }
-}
+
+    @Test
+    void runAll() {
+        AtomicBoolean ran1 = new AtomicBoolean(false);
+        AtomicBoolean ran2 = new AtomicBoolean(false);
+        Runnable runnable1 = () -> ran1.set(true);
+        Runnable runnable2 = () -> ran2.set(true);
+
+        CompletableFuture<?> runCF = CompletableFutureKit.runAll(List.of(runnable1, runnable2));
+        runCF.join();
+        assertThat(ran1.get(), equalTo(true));
+        assertThat(ran2.get(), equalTo(true));
+
+        ran1.set(false);
+        ran2.set(false);
+        Runnable runnable3 = () -> {
+            throw new RuntimeException("BANG");
+        };
+
+        runCF = CompletableFutureKit.runAll(List.of(runnable1, runnable2, runnable3));
+
+        CompletionException completionException = assertThrows(CompletionException.class, runCF::join);
+        assertThat(ran1.get(), equalTo(true));
+        assertThat(ran2.get(), equalTo(true));
+        assertThat(completionException.getCause().getMessage(), equalTo("BANG"));
+    }
+
+}

src/test/java/org/dataloader/scheduler/BatchLoaderSchedulerTest.java

@@ -3,6 +3,7 @@
 import org.dataloader.BatchLoaderEnvironment;
 import org.dataloader.DataLoader;
 import org.dataloader.DataLoaderOptions;
+import org.dataloader.impl.CompletableFutureKit;
 import org.junit.jupiter.api.Test;
 
 import java.util.List;
@@ -11,6 +12,7 @@
 import java.util.concurrent.CompletionStage;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Function;
+import java.util.stream.Collectors;
 
 import static org.awaitility.Awaitility.await;
 import static org.dataloader.DataLoaderFactory.newDataLoader;
@@ -188,6 +190,7 @@ public <K> void scheduleBatchPublisher(ScheduledBatchPublisherCall scheduledCall
     void can_schedule_cf_completion() {
 
         AtomicBoolean useThreading = new AtomicBoolean(false);
+        AtomicBoolean parallelCompletion = new AtomicBoolean(false);
         BatchLoaderScheduler scheduler = new BatchLoaderScheduler() {
             @Override
             public <K, V> CompletionStage<List<V>> scheduleBatchLoader(ScheduledBatchLoaderCall<V> scheduledCall, List<K> keys, BatchLoaderEnvironment environment) {
@@ -205,12 +208,17 @@ public <K> void scheduleBatchPublisher(ScheduledBatchPublisherCall scheduledCall
             }
 
             @Override
-            public <K> CompletionStage<Void> scheduleCompletion(Runnable completeValuesRunnable, List<K> keys, BatchLoaderEnvironment environment) {
+            public <K> CompletionStage<Void> scheduleCompletion(List<Runnable> completeValueRunnables, List<K> keys, BatchLoaderEnvironment environment) {
                 if (useThreading.get()) {
                     snooze(500);
-                    return CompletableFuture.runAsync(completeValuesRunnable);
+                    if (!parallelCompletion.get()) {
+                        return CompletableFutureKit.runAll(completeValueRunnables);
+                    } else {
+                        return CompletableFuture.allOf(completeValueRunnables.stream()
+                                .map(CompletableFuture::runAsync).toArray(CompletableFuture[]::new));
+                    }
                 } else {
-                    return BatchLoaderScheduler.super.scheduleCompletion(completeValuesRunnable, keys, environment);
+                    return BatchLoaderScheduler.super.scheduleCompletion(completeValueRunnables, keys, environment);
                 }
             }
         };
@@ -227,7 +235,20 @@ public <K> CompletionStage<Void> scheduleCompletion(Runnable completeValuesRunna
         assertThat(cf1.join(), equalTo(1));
         assertThat(cf2.join(), equalTo(2));
 
-        // switch mode to threading mdoe
+        // switch mode to threading mode
+
+        useThreading.set(true);
+
+        cf1 = identityLoader.load(10);
+        cf2 = identityLoader.load(20);
+        dispatchCF = identityLoader.dispatch();
+
+
+        await().until(dispatchCF::isDone);
+        assertThat(cf1.join(), equalTo(10));
+        assertThat(cf2.join(), equalTo(20));
+
+        // switch mode to parallel execution
 
         useThreading.set(true);