Added batch max size plus exceptions in the list support

bbakerman · bbakerman · commit 5f7ba3685df0 · 2017-07-04T20:53:15.000+10:00
diff --git a/README.md b/README.md
@@ -196,22 +196,21 @@ and there are also gains to this different mode of operation:
 However, with batch execution control comes responsibility! If you forget to make the call to `dispatch()` then the futures
 in the load request queue will never be batched, and thus _will never complete_! So be careful when crafting your loader designs.
 
-### Error object is not a thing in Java
+### Error object is not a thing in a type safe Java world
 
 In the reference JS implementation if the batch loader returns an `Error` object back then the `loadKey()` promise is rejected
-with that error.  This allows fine grain (per object in the list)  sets of error.  If I ask for keys A,B,C and B errors out the promise
-for B can contain a specific eror. 
+with that error.  This allows fine grain (per object in the list) sets of error.  If I ask for keys A,B,C and B errors out the promise
+for B can contain a specific error. 
 
-This is not (easily) possible in a Java implementation
+This is not quite as neat in a Java implementation
 
-A batch loader function is defined as `BatchLoader<K, V>` meaning for a key of type `K` it returns a value of type `V`.  It can just return
-some `Exception` as an object of type `V` since Java is type safe.
-  
-We could have made it return an `List<Either<V,Exception>>` but this greatly complicates the method signatures and so this aspect was not ported.
-   
-Instead when the batch loader function is called, it returns a promise of `List<V>` and if that promise fails in aggregate then that is what
-is reported back.   
+A batch loader function is defined as `BatchLoader<K, V>` meaning for a key of type `K` it returns a value of type `V`.  
+
+It cant just return some `Exception` as an object of type `V` since Java is type safe.
 
+You in order for a batch loader function to return an `Exception` it must be declared as `BatchLoader<K, Object>` which
+allows both values and exceptions to be returned .  Some type safety is lost in this case if you want
+to use the mix of exceptions and values pattern.
 
 ## Let's get started!
 
diff --git a/src/main/java/org/dataloader/DataLoader.java b/src/main/java/org/dataloader/DataLoader.java
@@ -19,14 +19,15 @@
 import org.dataloader.impl.CompletableFutureKit;
 
 import java.util.ArrayList;
-import java.util.Collections;
+import java.util.Collection;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 import java.util.stream.Collectors;
 
 import static java.util.Collections.emptyList;
+import static java.util.Collections.singletonList;
 import static org.dataloader.impl.Assertions.assertState;
 import static org.dataloader.impl.Assertions.nonNull;
 
@@ -98,8 +99,10 @@ private CacheMap<Object, CompletableFuture<V>> determineCacheMap(DataLoaderOptio
      */
     public CompletableFuture<V> load(K key) {
         Object cacheKey = getCacheKey(nonNull(key));
-        if (loaderOptions.cachingEnabled() && futureCache.containsKey(cacheKey)) {
-            return futureCache.get(cacheKey);
+        synchronized (futureCache) {
+            if (loaderOptions.cachingEnabled() && futureCache.containsKey(cacheKey)) {
+                return futureCache.get(cacheKey);
+            }
         }
 
         CompletableFuture<V> future = new CompletableFuture<>();
@@ -108,15 +111,17 @@ public CompletableFuture<V> load(K key) {
                 loaderQueue.put(key, future);
             }
         } else {
-            // immediate execution of batch function (but not promise itself)
+            // immediate execution of batch function
             CompletableFuture<List<V>> batchedLoad = batchLoadFunction
-                    .load(Collections.singletonList(key))
+                    .load(singletonList(key))
                     .toCompletableFuture();
             future = batchedLoad
                     .thenApply(list -> list.get(0));
         }
         if (loaderOptions.cachingEnabled()) {
-            futureCache.set(cacheKey, future);
+            synchronized (futureCache) {
+                futureCache.set(cacheKey, future);
+            }
         }
         return future;
     }
@@ -177,25 +182,64 @@ public CompletableFuture<List<V>> dispatch() {
         // the previously cached future objects that the client already has been given
         // via calls to load("foo") and loadMany(["foo","bar"])
         //
+        int maxBatchSize = loaderOptions.maxBatchSize();
+        if (maxBatchSize > 0 && maxBatchSize < keys.size()) {
+            return sliceIntoBatchesOfBatches(keys, queuedFutures, maxBatchSize);
+        } else {
+            return dispatchQueueBatch(keys, queuedFutures);
+        }
+    }
+
+    private CompletableFuture<List<V>> sliceIntoBatchesOfBatches(List<K> keys, List<CompletableFuture<V>> queuedFutures, int maxBatchSize) {
+        // the number of keys is > than what the batch loader function can accept
+        // so make multiple calls to the loader
+        List<CompletableFuture<List<V>>> allBatches = new ArrayList<>();
+        int len = keys.size();
+        int batchCount = (int) Math.ceil(len / (double) maxBatchSize);
+        for (int i = 0; i < batchCount; i++) {
+
+            int fromIndex = i * maxBatchSize;
+            int toIndex = Math.min((i + 1) * maxBatchSize, len);
+
+            List<K> subKeys = keys.subList(fromIndex, toIndex);
+            List<CompletableFuture<V>> subFutures = queuedFutures.subList(fromIndex, toIndex);
+
+            allBatches.add(dispatchQueueBatch(subKeys, subFutures));
+        }
+        //
+        // now reassemble all the futures into one that is the complete set of results
+        return CompletableFuture.allOf(allBatches.toArray(new CompletableFuture[allBatches.size()]))
+                .thenApply(v -> allBatches.stream()
+                        .map(CompletableFuture::join)
+                        .flatMap(Collection::stream)
+                        .collect(Collectors.toList()));
+    }
+
+    private CompletableFuture<List<V>> dispatchQueueBatch(List<K> keys, List<CompletableFuture<V>> queuedFutures) {
         return batchLoadFunction.load(keys)
                 .toCompletableFuture()
                 .thenApply(values -> {
                     assertState(keys.size() == values.size(), "The size of the promised values MUST be the same size as the key list");
 
                     for (int idx = 0; idx < queuedFutures.size(); idx++) {
-                        V value = values.get(idx);
+                        Object value = values.get(idx);
                         CompletableFuture<V> future = queuedFutures.get(idx);
-                        future.complete(value);
+                        if (value instanceof Throwable) {
+                            future.completeExceptionally((Throwable) value);
+                        } else {
+                            @SuppressWarnings("unchecked")
+                            V val = (V) value;
+                            future.complete(val);
+                        }
                     }
                     return values;
                 }).exceptionally(ex -> {
                     for (int idx = 0; idx < queuedFutures.size(); idx++) {
                         K key = keys.get(idx);
                         CompletableFuture<V> future = queuedFutures.get(idx);
-                        // clear any cached view of this key
-                        futureCache.delete(key);
                         future.completeExceptionally(ex);
-
+                        // clear any cached view of this key
+                        clear(key);
                     }
                     return emptyList();
                 });
@@ -242,7 +286,9 @@ public int dispatchDepth() {
      */
     public DataLoader<K, V> clear(K key) {
         Object cacheKey = getCacheKey(key);
-        futureCache.delete(cacheKey);
+        synchronized (futureCache) {
+            futureCache.delete(cacheKey);
+        }
         return this;
     }
 
@@ -252,7 +298,9 @@ public DataLoader<K, V> clear(K key) {
      * @return the data loader for fluent coding
      */
     public DataLoader<K, V> clearAll() {
-        futureCache.clear();
+        synchronized (futureCache) {
+            futureCache.clear();
+        }
         return this;
     }
 
@@ -266,8 +314,10 @@ public DataLoader<K, V> clearAll() {
      */
     public DataLoader<K, V> prime(K key, V value) {
         Object cacheKey = getCacheKey(key);
-        if (!futureCache.containsKey(cacheKey)) {
-            futureCache.set(cacheKey, CompletableFuture.completedFuture(value));
+        synchronized (futureCache) {
+            if (!futureCache.containsKey(cacheKey)) {
+                futureCache.set(cacheKey, CompletableFuture.completedFuture(value));
+            }
         }
         return this;
     }
diff --git a/src/main/java/org/dataloader/DataLoaderOptions.java b/src/main/java/org/dataloader/DataLoaderOptions.java
@@ -31,13 +31,15 @@ public class DataLoaderOptions {
     private boolean cachingEnabled;
     private CacheKey cacheKeyFunction;
     private CacheMap cacheMap;
+    private int maxBatchSize;
 
     /**
      * Creates a new data loader options with default settings.
      */
     public DataLoaderOptions() {
         batchingEnabled = true;
         cachingEnabled = true;
+        maxBatchSize = -1;
     }
 
     /**
@@ -51,6 +53,7 @@ public DataLoaderOptions(DataLoaderOptions other) {
         this.cachingEnabled = other.cachingEnabled;
         this.cacheKeyFunction = other.cacheKeyFunction;
         this.cacheMap = other.cacheMap;
+        this.maxBatchSize = other.maxBatchSize;
     }
 
     public static DataLoaderOptions create() {
@@ -144,4 +147,27 @@ public DataLoaderOptions setCacheMap(CacheMap cacheMap) {
         this.cacheMap = cacheMap;
         return this;
     }
+
+    /**
+     * Gets the maximum number of keys that will be presented to the {@link BatchLoader} function
+     * before they are split into multiple class
+     *
+     * @return the maximum batch size or -1 if there is no limit
+     */
+    public int maxBatchSize() {
+        return maxBatchSize;
+    }
+
+    /**
+     * Sets the maximum number of keys that will be presented to the {@link BatchLoader} function
+     * before they are split into multiple class
+     *
+     * @param maxBatchSize the maximum batch size
+     *
+     * @return the data loader options for fluent coding
+     */
+    public DataLoaderOptions setMaxBatchSize(int maxBatchSize) {
+        this.maxBatchSize = maxBatchSize;
+        return this;
+    }
 }
diff --git a/src/test/java/org/dataloader/DataLoaderTest.java b/src/test/java/org/dataloader/DataLoaderTest.java
