1717package org .dataloader ;
1818
1919import org .dataloader .impl .CompletableFutureKit ;
20+ import org .dataloader .stats .Statistics ;
21+ import org .dataloader .stats .StatisticsCollector ;
2022
2123import java .util .ArrayList ;
2224import java .util .Collection ;
2325import java .util .LinkedHashMap ;
2426import java .util .List ;
2527import java .util .Map ;
2628import java .util .concurrent .CompletableFuture ;
29+ import java .util .concurrent .CompletionStage ;
2730import java .util .stream .Collectors ;
2831
2932import static java .util .Collections .emptyList ;
@@ -64,6 +67,7 @@ public class DataLoader<K, V> {
6467 private final DataLoaderOptions loaderOptions ;
6568 private final CacheMap <Object , CompletableFuture <V >> futureCache ;
6669 private final Map <K , CompletableFuture <V >> loaderQueue ;
70+ private final StatisticsCollector stats ;
6771
6872 /**
6973 * Creates new DataLoader with the specified batch loader function and default options
@@ -153,6 +157,7 @@ public DataLoader(BatchLoader<K, V> batchLoadFunction, DataLoaderOptions options
153157 this .futureCache = determineCacheMap (loaderOptions );
154158 // order of keys matter in data loader
155159 this .loaderQueue = new LinkedHashMap <>();
160+ this .stats = nonNull (this .loaderOptions .getStatisticsCollector ());
156161 }
157162
158163 @ SuppressWarnings ("unchecked" )
@@ -173,8 +178,11 @@ private CacheMap<Object, CompletableFuture<V>> determineCacheMap(DataLoaderOptio
173178 */
174179 public CompletableFuture <V > load (K key ) {
175180 Object cacheKey = getCacheKey (nonNull (key ));
181+ stats .incrementLoadCount ();
182+
176183 synchronized (futureCache ) {
177184 if (loaderOptions .cachingEnabled () && futureCache .containsKey (cacheKey )) {
185+ stats .incrementCacheHitCount ();
178186 return futureCache .get (cacheKey );
179187 }
180188 }
@@ -185,6 +193,7 @@ public CompletableFuture<V> load(K key) {
185193 loaderQueue .put (key , future );
186194 }
187195 } else {
196+ stats .incrementBatchLoadCountBy (1 );
188197 // immediate execution of batch function
189198 CompletableFuture <List <V >> batchedLoad = batchLoadFunction
190199 .load (singletonList (key ))
@@ -291,7 +300,14 @@ private CompletableFuture<List<V>> sliceIntoBatchesOfBatches(List<K> keys, List<
291300
292301 @ SuppressWarnings ("unchecked" )
293302 private CompletableFuture <List <V >> dispatchQueueBatch (List <K > keys , List <CompletableFuture <V >> queuedFutures ) {
294- return batchLoadFunction .load (keys )
303+ stats .incrementBatchLoadCountBy (keys .size ());
304+ CompletionStage <List <V >> batchLoad ;
305+ try {
306+ batchLoad = nonNull (batchLoadFunction .load (keys ), "Your batch loader function MUST return a non null CompletionStage promise" );
307+ } catch (Exception e ) {
308+ batchLoad = CompletableFutureKit .failedFuture (e );
309+ }
310+ return batchLoad
295311 .toCompletableFuture ()
296312 .thenApply (values -> {
297313 assertState (keys .size () == values .size (), "The size of the promised values MUST be the same size as the key list" );
@@ -300,20 +316,28 @@ private CompletableFuture<List<V>> dispatchQueueBatch(List<K> keys, List<Complet
300316 Object value = values .get (idx );
301317 CompletableFuture <V > future = queuedFutures .get (idx );
302318 if (value instanceof Throwable ) {
319+ stats .incrementLoadErrorCount ();
303320 future .completeExceptionally ((Throwable ) value );
304321 // we don't clear the cached view of this entry to avoid
305322 // frequently loading the same error
306323 } else if (value instanceof Try ) {
307324 // we allow the batch loader to return a Try so we can better represent a computation
308325 // that might have worked or not.
309- handleTry ((Try <V >) value , future );
326+ Try <V > tryValue = (Try <V >) value ;
327+ if (tryValue .isSuccess ()) {
328+ future .complete (tryValue .get ());
329+ } else {
330+ stats .incrementLoadErrorCount ();
331+ future .completeExceptionally (tryValue .getThrowable ());
332+ }
310333 } else {
311334 V val = (V ) value ;
312335 future .complete (val );
313336 }
314337 }
315338 return values ;
316339 }).exceptionally (ex -> {
340+ stats .incrementBatchLoadExceptionCount ();
317341 for (int idx = 0 ; idx < queuedFutures .size (); idx ++) {
318342 K key = keys .get (idx );
319343 CompletableFuture <V > future = queuedFutures .get (idx );
@@ -325,14 +349,6 @@ private CompletableFuture<List<V>> dispatchQueueBatch(List<K> keys, List<Complet
325349 });
326350 }
327351
328- private void handleTry (Try <V > vTry , CompletableFuture <V > future ) {
329- if (vTry .isSuccess ()) {
330- future .complete (vTry .get ());
331- } else {
332- future .completeExceptionally (vTry .getThrowable ());
333- }
334- }
335-
336352 /**
337353 * Normally {@link #dispatch()} is an asynchronous operation but this version will 'join' on the
338354 * results if dispatch and wait for them to complete. If the {@link CompletableFuture} callbacks make more
@@ -441,4 +457,15 @@ public Object getCacheKey(K key) {
441457 return loaderOptions .cacheKeyFunction ().isPresent () ?
442458 loaderOptions .cacheKeyFunction ().get ().getKey (key ) : key ;
443459 }
460+
461+ /**
462+ * Gets the statistics associated with this data loader. These will have been gather via
463+ * the {@link org.dataloader.stats.StatisticsCollector} passed in via {@link DataLoaderOptions#getStatisticsCollector()}
464+ *
465+ * @return statistics for this data loader
466+ */
467+ public Statistics getStatistics () {
468+ return stats .getStatistics ();
469+ }
470+
444471}
0 commit comments