@@ -52,6 +52,7 @@ public class DataLoader<K, V> {
5252 private final DataLoaderOptions loaderOptions ;
5353 private final CacheMap <Object , Future <V >> futureCache ;
5454 private final LinkedHashMap <K , Future <V >> loaderQueue ;
55+ private final LinkedHashMap <CompositeFuture , LinkedHashMap <K , Future <V >>> dispatchedQueues ;
5556
5657 /**
5758 * Creates a new data loader with the provided batch load function, and default options.
@@ -75,6 +76,7 @@ public DataLoader(BatchLoader<K> batchLoadFunction, DataLoaderOptions options) {
7576 this .loaderOptions = options == null ? new DataLoaderOptions () : options ;
7677 this .futureCache = loaderOptions .cacheMap ().isPresent () ? (CacheMap <Object , Future <V >>) loaderOptions .cacheMap ().get () : CacheMap .simpleMap ();
7778 this .loaderQueue = new LinkedHashMap <>();
79+ this .dispatchedQueues = new LinkedHashMap <>();
7880 }
7981
8082 /**
@@ -138,18 +140,20 @@ public CompositeFuture dispatch() {
138140 return CompositeFuture .join (Collections .emptyList ());
139141 }
140142 CompositeFuture batch = batchLoadFunction .load (loaderQueue .keySet ());
143+ dispatchedQueues .put (batch , new LinkedHashMap <>(loaderQueue ));
141144 batch .setHandler (rh -> {
142145 AtomicInteger index = new AtomicInteger (0 );
143- loaderQueue .forEach ((key , future ) -> {
146+ dispatchedQueues . get ( batch ) .forEach ((key , future ) -> {
144147 if (batch .succeeded (index .get ())) {
145148 future .complete (batch .resultAt (index .get ()));
146149 } else {
147150 future .fail (batch .cause (index .get ()));
148151 }
149152 index .incrementAndGet ();
150153 });
151- loaderQueue . clear ( );
154+ dispatchedQueues . remove ( batch );
152155 });
156+ loaderQueue .clear ();
153157 return batch ;
154158 }
155159
0 commit comments