Updating the PR followup work

Author: bbakerman

build.gradle

@@ -61,5 +61,5 @@ dependencies {
 
 task wrapper(type: Wrapper) {
     gradleVersion = '4.0'
-    distributionUrl = "http://services.gradle.org/distributions/gradle-4.0-all.zip"
+    distributionUrl = "https://services.gradle.org/distributions/gradle-4.0-all.zip"
 }

gradle/publishing.gradle

@@ -4,7 +4,7 @@ apply plugin: 'com.jfrog.bintray'
 publishing {
     repositories {
         maven {
-            url 'http://dl.bintray.com/bbakerman/java-dataloader'
+            url 'https://dl.bintray.com/bbakerman/java-dataloader'
         }
     }
 }
@@ -52,7 +52,7 @@ publishing {
                     licenses {
                         license {
                             name 'The Apache Software License, Version 2.0'
-                            url 'http://www.apache.org/licenses/LICENSE-2.0.txt'
+                            url 'https://www.apache.org/licenses/LICENSE-2.0.txt'
                             distribution 'repo'
                         }
                     }

gradle/wrapper/gradle-wrapper.properties

@@ -3,4 +3,4 @@ distributionBase=GRADLE_USER_HOME
 distributionPath=wrapper/dists
 zipStoreBase=GRADLE_USER_HOME
 zipStorePath=wrapper/dists
-distributionUrl=http\://services.gradle.org/distributions/gradle-4.0-all.zip
+distributionUrl=https\://services.gradle.org/distributions/gradle-4.0-all.zip

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

@@ -470,7 +470,7 @@ 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;
+        return helper.dispatch().getPromisedResults();
     }
 
     /**
@@ -480,9 +480,9 @@ public CompletableFuture<List<V>> dispatch() {
      * 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.
+     * @return the promise of the queued load requests and the number of keys dispatched.
      */
-    public DataLoaderHelper.DispatchResult<V> dispatchWithCounts() {
+    public DispatchResult<V> dispatchWithCounts() {
         return helper.dispatch();
     }
 

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

@@ -134,15 +134,6 @@ Object getCacheKey(K key) {
                 loaderOptions.cacheKeyFunction().get().getKey(key) : key;
     }
 
-    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();
         //

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

@@ -118,8 +118,8 @@ public void dispatchAll() {
      */
     public int dispatchAllWithCount() {
         int sum = 0;
-        for (DataLoader dataLoader : getDataLoaders()) {
-            sum += dataLoader.dispatchWithCounts().totalEntriesHandled;
+        for (DataLoader<?,?> dataLoader : getDataLoaders()) {
+            sum += dataLoader.dispatchWithCounts().getKeysCount();
         }
         return sum;
     }

src/main/java/org/dataloader/DispatchResult.java

@@ -0,0 +1,29 @@
+package org.dataloader;
+
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+
+/**
+ * When a DataLoader is dispatched this object holds the promised results and also the count of key asked for
+ * via methods like {@link org.dataloader.DataLoader#load(Object)} or {@link org.dataloader.DataLoader#loadMany(java.util.List)}
+ *
+ * @param <T> for two
+ */
+
+public class DispatchResult<T> {
+    private final CompletableFuture<List<T>> futureList;
+    private final int keysCount;
+
+    public DispatchResult(CompletableFuture<List<T>> futureList, int keysCount) {
+        this.futureList = futureList;
+        this.keysCount = keysCount;
+    }
+
+    public CompletableFuture<List<T>> getPromisedResults() {
+        return futureList;
+    }
+
+    public int getKeysCount() {
+        return keysCount;
+    }
+}

src/test/java/org/dataloader/DataLoaderRegistryTest.java

@@ -1,9 +1,10 @@
 package org.dataloader;
 
-import java.util.concurrent.CompletableFuture;
 import org.dataloader.stats.Statistics;
 import org.junit.Test;
 
+import java.util.concurrent.CompletableFuture;
+
 import static java.util.Arrays.asList;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.hasItems;
@@ -131,4 +132,28 @@ public void computeIfAbsent_returns_an_existing_data_loader_if_there_was_a_value
         assertThat(registry.getDataLoaders(), hasItems(dlA));
     }
 
+    @Test
+    public void dispatch_counts_are_maintained() {
+
+        DataLoaderRegistry registry = new DataLoaderRegistry();
+
+        DataLoader<Object, Object> dlA = new DataLoader<>(identityBatchLoader);
+        DataLoader<Object, Object> dlB = new DataLoader<>(identityBatchLoader);
+
+        registry.register("a", dlA);
+        registry.register("b", dlB);
+
+        dlA.load("av1");
+        dlA.load("av2");
+        dlB.load("bv1");
+        dlB.load("bv2");
+
+        int dispatchDepth = registry.dispatchDepth();
+        assertThat(dispatchDepth, equalTo(4));
+
+        int dispatchedCount = registry.dispatchAllWithCount();
+        dispatchDepth = registry.dispatchDepth();
+        assertThat(dispatchedCount, equalTo(4));
+        assertThat(dispatchDepth, equalTo(0));
+    }
 }

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

@@ -110,10 +110,11 @@ public void should_Return_zero_entries_dispatched_when_no_keys_supplied() {
             assertThat(promisedValues.size(), is(0));
             success.set(true);
         });
-        DataLoaderHelper.DispatchResult dispatchResult = identityLoader.dispatchWithCounts();
+        DispatchResult dispatchResult = identityLoader.dispatchWithCounts();
         await().untilAtomic(success, is(true));
-        assertThat(dispatchResult.totalEntriesHandled, equalTo(0));
+        assertThat(dispatchResult.getKeysCount(), equalTo(0));
     }
+
     @Test
     public void should_Batch_multiple_requests() throws ExecutionException, InterruptedException {
         List<Collection<Integer>> loadCalls = new ArrayList<>();
@@ -126,19 +127,22 @@ 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> future1a = identityLoader.load(1);
         CompletableFuture<Integer> future2 = identityLoader.load(2);
-        DataLoaderHelper.DispatchResult dispatchResult = identityLoader.dispatchWithCounts();
+        DispatchResult<?> dispatchResult = identityLoader.dispatchWithCounts();
 
         await().until(() -> future1.isDone() && future2.isDone());
-        assertThat(dispatchResult.totalEntriesHandled, equalTo(2));
+        assertThat(dispatchResult.getKeysCount(), equalTo(2)); // its two because its the number dispatched (by key) not the load calls
+        assertThat(dispatchResult.getPromisedResults().isDone(), equalTo(true));
     }
 
     @Test