44import org .dataloader .annotations .Internal ;
55import org .dataloader .impl .CompletableFutureKit ;
66import org .dataloader .stats .StatisticsCollector ;
7+ import org .dataloader .stats .context .IncrementBatchLoadCountByStatisticsContext ;
8+ import org .dataloader .stats .context .IncrementBatchLoadExceptionCountStatisticsContext ;
9+ import org .dataloader .stats .context .IncrementCacheHitCountStatisticsContext ;
10+ import org .dataloader .stats .context .IncrementLoadCountStatisticsContext ;
11+ import org .dataloader .stats .context .IncrementLoadErrorCountStatisticsContext ;
712
813import java .time .Clock ;
914import java .time .Instant ;
@@ -105,7 +110,7 @@ Optional<CompletableFuture<V>> getIfPresent(K key) {
105110 if (cachingEnabled ) {
106111 Object cacheKey = getCacheKey (nonNull (key ));
107112 if (futureCache .containsKey (cacheKey )) {
108- stats .incrementCacheHitCount ();
113+ stats .incrementCacheHitCount (new IncrementCacheHitCountStatisticsContext <>( key ) );
109114 return Optional .of (futureCache .get (cacheKey ));
110115 }
111116 }
@@ -132,7 +137,7 @@ CompletableFuture<V> load(K key, Object loadContext) {
132137 boolean batchingEnabled = loaderOptions .batchingEnabled ();
133138 boolean cachingEnabled = loaderOptions .cachingEnabled ();
134139
135- stats .incrementLoadCount ();
140+ stats .incrementLoadCount (new IncrementLoadCountStatisticsContext <>( key , loadContext ) );
136141
137142 if (cachingEnabled ) {
138143 return loadFromCache (key , loadContext , batchingEnabled );
@@ -223,18 +228,20 @@ private CompletableFuture<List<V>> sliceIntoBatchesOfBatches(List<K> keys, List<
223228
224229 @ SuppressWarnings ("unchecked" )
225230 private CompletableFuture <List <V >> dispatchQueueBatch (List <K > keys , List <Object > callContexts , List <CompletableFuture <V >> queuedFutures ) {
226- stats .incrementBatchLoadCountBy (keys .size ());
231+ stats .incrementBatchLoadCountBy (keys .size (), new IncrementBatchLoadCountByStatisticsContext <>( keys , callContexts ) );
227232 CompletableFuture <List <V >> batchLoad = invokeLoader (keys , callContexts , loaderOptions .cachingEnabled ());
228233 return batchLoad
229234 .thenApply (values -> {
230235 assertResultSize (keys , values );
231236
232237 List <K > clearCacheKeys = new ArrayList <>();
233238 for (int idx = 0 ; idx < queuedFutures .size (); idx ++) {
239+ K key = keys .get (idx );
234240 V value = values .get (idx );
241+ Object callContext = callContexts .get (idx );
235242 CompletableFuture <V > future = queuedFutures .get (idx );
236243 if (value instanceof Throwable ) {
237- stats .incrementLoadErrorCount ();
244+ stats .incrementLoadErrorCount (new IncrementLoadErrorCountStatisticsContext <>( key , callContext ) );
238245 future .completeExceptionally ((Throwable ) value );
239246 clearCacheKeys .add (keys .get (idx ));
240247 } else if (value instanceof Try ) {
@@ -244,7 +251,7 @@ private CompletableFuture<List<V>> dispatchQueueBatch(List<K> keys, List<Object>
244251 if (tryValue .isSuccess ()) {
245252 future .complete (tryValue .get ());
246253 } else {
247- stats .incrementLoadErrorCount ();
254+ stats .incrementLoadErrorCount (new IncrementLoadErrorCountStatisticsContext <>( key , callContext ) );
248255 future .completeExceptionally (tryValue .getThrowable ());
249256 clearCacheKeys .add (keys .get (idx ));
250257 }
@@ -255,7 +262,7 @@ private CompletableFuture<List<V>> dispatchQueueBatch(List<K> keys, List<Object>
255262 possiblyClearCacheEntriesOnExceptions (clearCacheKeys );
256263 return values ;
257264 }).exceptionally (ex -> {
258- stats .incrementBatchLoadExceptionCount ();
265+ stats .incrementBatchLoadExceptionCount (new IncrementBatchLoadExceptionCountStatisticsContext <>( keys , callContexts ) );
259266 if (ex instanceof CompletionException ) {
260267 ex = ex .getCause ();
261268 }
@@ -294,7 +301,7 @@ private CompletableFuture<V> loadFromCache(K key, Object loadContext, boolean ba
294301
295302 if (futureCache .containsKey (cacheKey )) {
296303 // We already have a promise for this key, no need to check value cache or queue up load
297- stats .incrementCacheHitCount ();
304+ stats .incrementCacheHitCount (new IncrementCacheHitCountStatisticsContext <>( key , loadContext ) );
298305 return futureCache .get (cacheKey );
299306 }
300307
@@ -310,7 +317,7 @@ private CompletableFuture<V> queueOrInvokeLoader(K key, Object loadContext, bool
310317 loaderQueue .add (new LoaderQueueEntry <>(key , loadCallFuture , loadContext ));
311318 return loadCallFuture ;
312319 } else {
313- stats .incrementBatchLoadCountBy (1 );
320+ stats .incrementBatchLoadCountBy (1 , new IncrementBatchLoadCountByStatisticsContext <>( key , loadContext ) );
314321 // immediate execution of batch function
315322 return invokeLoaderImmediately (key , loadContext , cachingEnabled );
316323 }
0 commit comments