Fix dispatch hanging when load() is called twice with the same key

Author: BURJA Lucian

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

@@ -20,11 +20,10 @@
 import org.dataloader.stats.Statistics;
 import org.dataloader.stats.StatisticsCollector;
 
+import java.util.AbstractMap.SimpleImmutableEntry;
 import java.util.ArrayList;
 import java.util.Collection;
-import java.util.LinkedHashMap;
 import java.util.List;
-import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
 import java.util.stream.Collectors;
@@ -66,7 +65,7 @@ public class DataLoader<K, V> {
     private final BatchLoader<K, V> batchLoadFunction;
     private final DataLoaderOptions loaderOptions;
     private final CacheMap<Object, CompletableFuture<V>> futureCache;
-    private final Map<K, CompletableFuture<V>> loaderQueue;
+    private final List<SimpleImmutableEntry<K, CompletableFuture<V>>> loaderQueue;
     private final StatisticsCollector stats;
 
     /**
@@ -156,7 +155,7 @@ public DataLoader(BatchLoader<K, V> batchLoadFunction, DataLoaderOptions options
         this.loaderOptions = options == null ? new DataLoaderOptions() : options;
         this.futureCache = determineCacheMap(loaderOptions);
         // order of keys matter in data loader
-        this.loaderQueue = new LinkedHashMap<>();
+        this.loaderQueue = new ArrayList<>();
         this.stats = nonNull(this.loaderOptions.getStatisticsCollector());
     }
 
@@ -192,7 +191,7 @@ public CompletableFuture<V> load(K key) {
         CompletableFuture<V> future = new CompletableFuture<>();
         if (loaderOptions.batchingEnabled()) {
             synchronized (loaderQueue) {
-                loaderQueue.put(key, future);
+                loaderQueue.add(new SimpleImmutableEntry<>(key, future));
             }
         } else {
             stats.incrementBatchLoadCountBy(1);
@@ -247,9 +246,9 @@ public CompletableFuture<List<V>> dispatch() {
         final List<K> keys = new ArrayList<>();
         final List<CompletableFuture<V>> queuedFutures = new ArrayList<>();
         synchronized (loaderQueue) {
-            loaderQueue.forEach((key, future) -> {
-                keys.add(key);
-                queuedFutures.add(future);
+            loaderQueue.forEach(entry -> {
+                keys.add(entry.getKey());
+                queuedFutures.add(entry.getValue());
             });
             loaderQueue.clear();
         }

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

@@ -575,6 +575,23 @@ public void should_Disable_caching() throws ExecutionException, InterruptedExcep
         assertThat(loadCalls, equalTo(asList(asList("A", "B"),
                 asList("A", "C"), asList("A", "B", "C"))));
     }
+    @Test
+    public void should_work_with_duplicate_keys() throws ExecutionException, InterruptedException {
+        List<Collection<String>> loadCalls = new ArrayList<>();
+        DataLoader<String, String> identityLoader =
+                idLoader(newOptions().setCachingEnabled(false), loadCalls);
+
+        CompletableFuture<String> future1 = identityLoader.load("A");
+        CompletableFuture<String> future2 = identityLoader.load("B");
+        CompletableFuture<String> future3 = identityLoader.load("A");
+        identityLoader.dispatch();
+
+        await().until(() -> future1.isDone() && future2.isDone() && future3.isDone());
+        assertThat(future1.get(), equalTo("A"));
+        assertThat(future2.get(), equalTo("B"));
+        assertThat(future3.get(), equalTo("A"));
+        assertThat(loadCalls, equalTo(singletonList(asList("A", "B", "A"))));
+    }
 
     // It is resilient to job queue ordering