2020import org .dataloader .stats .Statistics ;
2121import org .dataloader .stats .StatisticsCollector ;
2222
23- import java .util .AbstractMap .SimpleImmutableEntry ;
2423import java .util .ArrayList ;
24+ import java .util .Collections ;
2525import java .util .List ;
2626import java .util .concurrent .CompletableFuture ;
27- import java .util .stream .Collectors ;
2827
2928import static org .dataloader .impl .Assertions .nonNull ;
3029
@@ -60,7 +59,6 @@ public class DataLoader<K, V> {
6059 private final DataLoaderHelper <K , V > helper ;
6160 private final DataLoaderOptions loaderOptions ;
6261 private final CacheMap <Object , CompletableFuture <V >> futureCache ;
63- private final List <SimpleImmutableEntry <K , CompletableFuture <V >>> loaderQueue ;
6462 private final StatisticsCollector stats ;
6563
6664 /**
@@ -237,6 +235,7 @@ public static <K, V> DataLoader<K, V> newMappedDataLoader(MappedBatchLoader<K, V
237235 * Using Try objects allows you to capture a value returned or an exception that might
238236 * have occurred trying to get a value. .
239237 * <p>
238+ *
240239 * @param batchLoadFunction the batch load function to use that uses {@link org.dataloader.Try} objects
241240 * @param <K> the key type
242241 * @param <V> the value type
@@ -357,10 +356,9 @@ private DataLoader(Object batchLoadFunction, DataLoaderOptions options) {
357356 this .loaderOptions = options == null ? new DataLoaderOptions () : options ;
358357 this .futureCache = determineCacheMap (loaderOptions );
359358 // order of keys matter in data loader
360- this .loaderQueue = new ArrayList <>();
361359 this .stats = nonNull (this .loaderOptions .getStatisticsCollector ());
362360
363- this .helper = new DataLoaderHelper <>(this , batchLoadFunction , this .loaderOptions , this .futureCache , this .loaderQueue , this . stats );
361+ this .helper = new DataLoaderHelper <>(this , batchLoadFunction , this .loaderOptions , this .futureCache , this .stats );
364362 }
365363
366364 @ SuppressWarnings ("unchecked" )
@@ -380,7 +378,26 @@ private CacheMap<Object, CompletableFuture<V>> determineCacheMap(DataLoaderOptio
380378 * @return the future of the value
381379 */
382380 public CompletableFuture <V > load (K key ) {
383- return helper .load (key );
381+ return load (key , null );
382+ }
383+
384+ /**
385+ * Requests to load the data with the specified key asynchronously, and returns a future of the resulting value.
386+ * <p>
387+ * If batching is enabled (the default), you'll have to call {@link DataLoader#dispatch()} at a later stage to
388+ * start batch execution. If you forget this call the future will never be completed (unless already completed,
389+ * and returned from cache).
390+ * <p>
391+ * The key context object may be useful in the batch loader interfaces such as {@link org.dataloader.BatchLoaderWithContext} or
392+ * {@link org.dataloader.MappedBatchLoaderWithContext} to help retrieve data.
393+ *
394+ * @param key the key to load
395+ * @param keyContext a context object that is specific to this key
396+ *
397+ * @return the future of the value
398+ */
399+ public CompletableFuture <V > load (K key , Object keyContext ) {
400+ return helper .load (key , keyContext );
384401 }
385402
386403 /**
@@ -396,11 +413,39 @@ public CompletableFuture<V> load(K key) {
396413 * @return the composite future of the list of values
397414 */
398415 public CompletableFuture <List <V >> loadMany (List <K > keys ) {
399- synchronized (this ) {
400- List <CompletableFuture <V >> collect = keys .stream ()
401- .map (this ::load )
402- .collect (Collectors .toList ());
416+ return loadMany (keys , Collections .emptyList ());
417+ }
403418
419+ /**
420+ * Requests to load the list of data provided by the specified keys asynchronously, and returns a composite future
421+ * of the resulting values.
422+ * <p>
423+ * If batching is enabled (the default), you'll have to call {@link DataLoader#dispatch()} at a later stage to
424+ * start batch execution. If you forget this call the future will never be completed (unless already completed,
425+ * and returned from cache).
426+ * <p>
427+ * The key context object may be useful in the batch loader interfaces such as {@link org.dataloader.BatchLoaderWithContext} or
428+ * {@link org.dataloader.MappedBatchLoaderWithContext} to help retrieve data.
429+ *
430+ * @param keys the list of keys to load
431+ * @param keyContexts the list of key calling context objects
432+ *
433+ * @return the composite future of the list of values
434+ */
435+ public CompletableFuture <List <V >> loadMany (List <K > keys , List <Object > keyContexts ) {
436+ nonNull (keys );
437+ nonNull (keyContexts );
438+
439+ synchronized (this ) {
440+ List <CompletableFuture <V >> collect = new ArrayList <>();
441+ for (int i = 0 ; i < keys .size (); i ++) {
442+ K key = keys .get (i );
443+ Object keyContext = null ;
444+ if (i < keyContexts .size ()) {
445+ keyContext = keyContexts .get (i );
446+ }
447+ collect .add (load (key , keyContext ));
448+ }
404449 return CompletableFutureKit .allOf (collect );
405450 }
406451 }
@@ -441,9 +486,7 @@ public List<V> dispatchAndJoin() {
441486 * @return the depth of the batched key loads that need to be dispatched
442487 */
443488 public int dispatchDepth () {
444- synchronized (this ) {
445- return loaderQueue .size ();
446- }
489+ return helper .dispatchDepth ();
447490 }
448491
449492
0 commit comments