@@ -52,6 +52,7 @@ public class DataLoader<K, V> {
52
52
private final DataLoaderOptions loaderOptions ;
53
53
private final CacheMap <Object , Future <V >> futureCache ;
54
54
private final LinkedHashMap <K , Future <V >> loaderQueue ;
55
+ private final LinkedHashMap <CompositeFuture , LinkedHashMap <K , Future <V >>> dispatchedQueues ;
55
56
56
57
/**
57
58
* 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) {
75
76
this .loaderOptions = options == null ? new DataLoaderOptions () : options ;
76
77
this .futureCache = loaderOptions .cacheMap ().isPresent () ? (CacheMap <Object , Future <V >>) loaderOptions .cacheMap ().get () : CacheMap .simpleMap ();
77
78
this .loaderQueue = new LinkedHashMap <>();
79
+ this .dispatchedQueues = new LinkedHashMap <>();
78
80
}
79
81
80
82
/**
@@ -138,18 +140,20 @@ public CompositeFuture dispatch() {
138
140
return CompositeFuture .join (Collections .emptyList ());
139
141
}
140
142
CompositeFuture batch = batchLoadFunction .load (loaderQueue .keySet ());
143
+ dispatchedQueues .put (batch , new LinkedHashMap <>(loaderQueue ));
141
144
batch .setHandler (rh -> {
142
145
AtomicInteger index = new AtomicInteger (0 );
143
- loaderQueue .forEach ((key , future ) -> {
146
+ dispatchedQueues . get ( batch ) .forEach ((key , future ) -> {
144
147
if (batch .succeeded (index .get ())) {
145
148
future .complete (batch .resultAt (index .get ()));
146
149
} else {
147
150
future .fail (batch .cause (index .get ()));
148
151
}
149
152
index .incrementAndGet ();
150
153
});
151
- loaderQueue . clear ( );
154
+ dispatchedQueues . remove ( batch );
152
155
});
156
+ loaderQueue .clear ();
153
157
return batch ;
154
158
}
155
159
0 commit comments