Allow scheduling of CF completion in parallel

Author: fxbonnet

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

@@ -325,40 +325,52 @@ private CompletableFuture<List<V>> dispatchQueueBatch(List<K> keys, List<Object>
         stats.incrementBatchLoadCountBy(keys.size(), new IncrementBatchLoadCountByStatisticsContext<>(keys, callContexts));
         CompletableFuture<List<V>> batchLoad = invokeLoader(keys, callContexts, queuedFutures, loaderOptions.cachingEnabled());
         return batchLoad
-                .thenApply(values -> {
+                .thenCompose(values -> {
                     assertResultSize(keys, values);
                     if (isPublisher() || isMappedPublisher()) {
                         // We have already completed the queued futures by the time the overall batchLoad future has completed.
-                        return values;
+                        return CompletableFuture.completedFuture(values);
                     }
 
                     List<K> clearCacheKeys = new ArrayList<>();
+                    var batchLoaderScheduler = loaderOptions.getBatchLoaderScheduler();
+                    CompletableFuture<Void>[] scheduledCompletions = new CompletableFuture[keys.size()];
                     for (int idx = 0; idx < queuedFutures.size(); idx++) {
                         K key = keys.get(idx);
                         V value = values.get(idx);
                         Object callContext = callContexts.get(idx);
                         CompletableFuture<V> future = queuedFutures.get(idx);
-                        if (value instanceof Throwable) {
-                            stats.incrementLoadErrorCount(new IncrementLoadErrorCountStatisticsContext<>(key, callContext));
-                            future.completeExceptionally((Throwable) value);
-                            clearCacheKeys.add(keys.get(idx));
-                        } 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.
-                            Try<V> tryValue = (Try<V>) value;
-                            if (tryValue.isSuccess()) {
-                                future.complete(tryValue.get());
-                            } else {
+                        Runnable completeValueRunnable = () -> {
+                            if (value instanceof Throwable) {
                                 stats.incrementLoadErrorCount(new IncrementLoadErrorCountStatisticsContext<>(key, callContext));
-                                future.completeExceptionally(tryValue.getThrowable());
-                                clearCacheKeys.add(keys.get(idx));
+                                future.completeExceptionally((Throwable) value);
+                                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.
+                                Try<V> tryValue = (Try<V>) value;
+                                if (tryValue.isSuccess()) {
+                                    future.complete(tryValue.get());
+                                } else {
+                                    stats.incrementLoadErrorCount(new IncrementLoadErrorCountStatisticsContext<>(key, callContext));
+                                    future.completeExceptionally(tryValue.getThrowable());
+                                    clearCacheKeys.add(key);
+                                }
+                            } else {
+                                future.complete(value);
                             }
+                        };
+                        if(batchLoaderScheduler != null) {
+                            scheduledCompletions[idx] = batchLoaderScheduler.scheduleCompletion(completeValueRunnable, key, value);
                         } else {
-                            future.complete(value);
+                            scheduledCompletions[idx] = CompletableFutureKit.run(completeValueRunnable);
                         }
                     }
-                    possiblyClearCacheEntriesOnExceptions(clearCacheKeys);
-                    return values;
+                    // Wait for all completions to return
+                    return allOf(scheduledCompletions).thenApply(ignored -> {
+                        possiblyClearCacheEntriesOnExceptions(clearCacheKeys);
+                        return values;
+                    });
                 }).exceptionally(ex -> {
                     stats.incrementBatchLoadExceptionCount(new IncrementBatchLoadExceptionCountStatisticsContext<>(keys, callContexts));
                     if (ex instanceof CompletionException) {

src/main/java/org/dataloader/impl/CompletableFutureKit.java

@@ -67,4 +67,20 @@ public static <K, V> CompletableFuture<Map<K, V>> allOf(Map<K, CompletableFuture
                         )
                 );
     }
+
+    /**
+     * 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();
+            return CompletableFuture.completedFuture(null);
+        } catch (Exception e) {
+            return CompletableFuture.failedFuture(e);
+        }
+    }
 }

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

@@ -7,9 +7,11 @@
 import org.dataloader.MappedBatchLoader;
 import org.dataloader.MappedBatchPublisher;
 import org.dataloader.BatchPublisher;
+import org.dataloader.impl.CompletableFutureKit;
 
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
 
 /**
@@ -92,4 +94,21 @@ interface ScheduledBatchPublisherCall {
      * @param <K>           the key type
      */
     <K> void scheduleBatchPublisher(ScheduledBatchPublisherCall scheduledCall, List<K> keys, BatchLoaderEnvironment environment);
+
+    /**
+     * 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.
+     * <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
+     */
+    default <K, V> CompletableFuture<Void> scheduleCompletion(Runnable completeValueRunnable, K key, V value) {
+        return CompletableFutureKit.run(completeValueRunnable);
+    }
 }