Merge pull request #1 from getmargin/feature/dataloader-counts

Allow dispatch caller to get stats back

Author: Keith Dreibelbis

src/main/java/org/dataloader/DataLoader.java

@@ -470,6 +470,19 @@ public CompletableFuture<List<V>> loadMany(List<K> keys, List<Object> keyContext
      * @return the promise of the queued load requests
      */
     public CompletableFuture<List<V>> dispatch() {
+        return helper.dispatch().futureList;
+    }
+
+    /**
+     * Dispatches the queued load requests to the batch execution function and returns both the promise of the result
+     * and the number of entries that were dispatched.
+     * <p>
+     * If batching is disabled, or there are no queued requests, then a succeeded promise with no entries dispatched is
+     * returned.
+     *
+     * @return the promise of the queued load requests and the number of entries dispatched.
+     */
+    public DataLoaderHelper.DispatchResult<V> dispatchWithCounts() {
         return helper.dispatch();
     }
 

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

@@ -134,7 +134,16 @@ Object getCacheKey(K key) {
                 loaderOptions.cacheKeyFunction().get().getKey(key) : key;
     }
 
-    CompletableFuture<List<V>> dispatch() {
+    public static class DispatchResult<X> {
+        public final CompletableFuture<List<X>> futureList;
+        public final int totalEntriesHandled;
+        public DispatchResult(CompletableFuture<List<X>> futureList, int totalEntriesHandled) {
+            this.futureList = futureList;
+            this.totalEntriesHandled = totalEntriesHandled;
+        }
+    }
+
+    DispatchResult<V> dispatch() {
         boolean batchingEnabled = loaderOptions.batchingEnabled();
         //
         // we copy the pre-loaded set of futures ready for dispatch
@@ -149,9 +158,10 @@ CompletableFuture<List<V>> dispatch() {
             });
             loaderQueue.clear();
         }
-        if (!batchingEnabled || keys.size() == 0) {
-            return CompletableFuture.completedFuture(emptyList());
+        if (!batchingEnabled || keys.isEmpty()) {
+            return new DispatchResult<V>(CompletableFuture.completedFuture(emptyList()), 0);
         }
+        final int totalEntriesHandled = keys.size();
         //
         // order of keys -> values matter in data loader hence the use of linked hash map
         //
@@ -164,11 +174,13 @@ CompletableFuture<List<V>> dispatch() {
         // via calls to load("foo") and loadMany(["foo","bar"])
         //
         int maxBatchSize = loaderOptions.maxBatchSize();
+        CompletableFuture<List<V>> futureList;
         if (maxBatchSize > 0 && maxBatchSize < keys.size()) {
-            return sliceIntoBatchesOfBatches(keys, queuedFutures, callContexts, maxBatchSize);
+            futureList = sliceIntoBatchesOfBatches(keys, queuedFutures, callContexts, maxBatchSize);
         } else {
-            return dispatchQueueBatch(keys, callContexts, queuedFutures);
+            futureList = dispatchQueueBatch(keys, callContexts, queuedFutures);
         }
+        return new DispatchResult<V>(futureList, totalEntriesHandled);
     }
 
     private CompletableFuture<List<V>> sliceIntoBatchesOfBatches(List<K> keys, List<CompletableFuture<V>> queuedFutures, List<Object> callContexts, int maxBatchSize) {

src/main/java/org/dataloader/DataLoaderRegistry.java

@@ -7,6 +7,7 @@
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.function.Function;
+
 import org.dataloader.stats.Statistics;
 
 /**
@@ -22,7 +23,6 @@ public class DataLoaderRegistry {
      *
      * @param key        the key to put the data loader under
      * @param dataLoader the data loader to register
-     *
      * @return this registry
      */
     public DataLoaderRegistry register(String key, DataLoader<?, ?> dataLoader) {
@@ -33,15 +33,14 @@ public DataLoaderRegistry register(String key, DataLoader<?, ?> dataLoader) {
     /**
      * Computes a data loader if absent or return it if it was
      * already registered at that key.
-     *
+     * <p>
      * Note: The entire method invocation is performed atomically,
      * so the function is applied at most once per key.
      *
-     * @param key the key of the data loader
+     * @param key             the key of the data loader
      * @param mappingFunction the function to compute a data loader
-     * @param <K> the type of keys
-     * @param <V> the type of values
-     *
+     * @param <K>             the type of keys
+     * @param <V>             the type of values
      * @return a data loader
      */
     @SuppressWarnings("unchecked")
@@ -55,7 +54,6 @@ public <K, V> DataLoader<K, V> computeIfAbsent(final String key,
      * and return a new combined registry
      *
      * @param registry the registry to combine into this registry
-     *
      * @return a new combined registry
      */
     public DataLoaderRegistry combine(DataLoaderRegistry registry) {
@@ -77,7 +75,6 @@ public DataLoaderRegistry combine(DataLoaderRegistry registry) {
      * This will unregister a new dataloader
      *
      * @param key the key of the data loader to unregister
-     *
      * @return this registry
      */
     public DataLoaderRegistry unregister(String key) {
@@ -91,7 +88,6 @@ public DataLoaderRegistry unregister(String key) {
      * @param key the key of the data loader
      * @param <K> the type of keys
      * @param <V> the type of values
-     *
      * @return a data loader or null if its not present
      */
     @SuppressWarnings("unchecked")
@@ -114,6 +110,32 @@ public void dispatchAll() {
         getDataLoaders().forEach(DataLoader::dispatch);
     }
 
+    /**
+     * Similar to {@link DataLoaderRegistry#dispatchAll()}, this calls {@link org.dataloader.DataLoader#dispatch()} on
+     * each of the registered {@link org.dataloader.DataLoader}s, but returns the number of dispatches.
+     *
+     * @return total number of entries that were dispatched from registered {@link org.dataloader.DataLoader}s.
+     */
+    public int dispatchAllWithCount() {
+        int sum = 0;
+        for (DataLoader dataLoader : getDataLoaders()) {
+            sum += dataLoader.dispatchWithCounts().totalEntriesHandled;
+        }
+        return sum;
+    }
+
+    /**
+     * @return The sum of all batched key loads that need to be dispatched from all registered
+     * {@link org.dataloader.DataLoader}s
+     */
+    public int dispatchDepth() {
+        int totalDispatchDepth = 0;
+        for (DataLoader dataLoader : getDataLoaders()) {
+            totalDispatchDepth += dataLoader.dispatchDepth();
+        }
+        return totalDispatchDepth;
+    }
+
     /**
      * @return a combined set of statistics for all data loaders in this registry presented
      * as the sum of all their statistics

src/test/java/org/dataloader/DataLoaderTest.java

@@ -101,6 +101,19 @@ public void should_Resolve_to_empty_list_when_no_keys_supplied() {
         assertThat(futureEmpty.join(), empty());
     }
 
+    @Test
+    public void should_Return_zero_entries_dispatched_when_no_keys_supplied() {
+        AtomicBoolean success = new AtomicBoolean();
+        DataLoader<Integer, Integer> identityLoader = new DataLoader<>(keysAsValues());
+        CompletableFuture<List<Integer>> futureEmpty = identityLoader.loadMany(emptyList());
+        futureEmpty.thenAccept(promisedValues -> {
+            assertThat(promisedValues.size(), is(0));
+            success.set(true);
+        });
+        DataLoaderHelper.DispatchResult dispatchResult = identityLoader.dispatchWithCounts();
+        await().untilAtomic(success, is(true));
+        assertThat(dispatchResult.totalEntriesHandled, equalTo(0));
+    }
     @Test
     public void should_Batch_multiple_requests() throws ExecutionException, InterruptedException {
         List<Collection<Integer>> loadCalls = new ArrayList<>();
@@ -113,7 +126,19 @@ public void should_Batch_multiple_requests() throws ExecutionException, Interrup
         await().until(() -> future1.isDone() && future2.isDone());
         assertThat(future1.get(), equalTo(1));
         assertThat(future2.get(), equalTo(2));
-        assertThat(loadCalls, equalTo(singletonList(asList(1, 2))));
+        assertThat(loadCalls, equalTo(singletonList(asList(1, 2))));    }
+
+    @Test
+    public void should_Return_number_of_batched_entries() throws ExecutionException, InterruptedException {
+        List<Collection<Integer>> loadCalls = new ArrayList<>();
+        DataLoader<Integer, Integer> identityLoader = idLoader(new DataLoaderOptions(), loadCalls);
+
+        CompletableFuture<Integer> future1 = identityLoader.load(1);
+        CompletableFuture<Integer> future2 = identityLoader.load(2);
+        DataLoaderHelper.DispatchResult dispatchResult = identityLoader.dispatchWithCounts();
+
+        await().until(() -> future1.isDone() && future2.isDone());
+        assertThat(dispatchResult.totalEntriesHandled, equalTo(2));
     }
 
     @Test