8
8
import java .time .Instant ;
9
9
import java .util .ArrayList ;
10
10
import java .util .Collection ;
11
- import java .util .LinkedHashMap ;
12
11
import java .util .LinkedHashSet ;
13
12
import java .util .List ;
14
13
import java .util .Map ;
@@ -335,45 +334,46 @@ CompletableFuture<List<V>> invokeLoader(List<K> keys, List<Object> keyContexts,
335
334
336
335
assertState (keys .size () == cachedValues .size (), () -> "The size of the cached values MUST be the same size as the key list" );
337
336
338
- LinkedHashMap <K , V > valuesInKeyOrder = new LinkedHashMap <>();
339
- List <K > cacheMissedKeys = new ArrayList <>();
340
- List <Object > cacheMissedContexts = new ArrayList <>();
337
+ // the following is NOT a Map because keys in data loader can repeat (by design)
338
+ // and hence "a","b","c","b" is a valid set of keys
339
+ List <Try <V >> valuesInKeyOrder = new ArrayList <>();
340
+ List <Integer > missedKeyIndexes = new ArrayList <>();
341
+ List <K > missedKeys = new ArrayList <>();
342
+ List <Object > missedKeyContexts = new ArrayList <>();
341
343
for (int i = 0 ; i < keys .size (); i ++) {
342
- K key = keys .get (i );
343
- Object keyContext = keyContexts .get (i );
344
344
Try <V > cacheGet = cachedValues .get (i );
345
- if (cacheGet .isSuccess ()) {
346
- valuesInKeyOrder .put (key , cacheGet .get ());
347
- } else {
348
- valuesInKeyOrder .put (key , null ); // an entry to be replaced later
349
- cacheMissedKeys .add (key );
350
- cacheMissedContexts .add (keyContext );
345
+ valuesInKeyOrder .add (cacheGet );
346
+ if (cacheGet .isFailure ()) {
347
+ missedKeyIndexes .add (i );
348
+ missedKeys .add (keys .get (i ));
349
+ missedKeyContexts .add (keyContexts .get (i ));
351
350
}
352
351
}
353
- if (cacheMissedKeys .isEmpty ()) {
352
+ if (missedKeys .isEmpty ()) {
354
353
//
355
354
// everything was cached
356
355
//
357
- return completedFuture (new ArrayList <>(valuesInKeyOrder .values ()));
356
+ List <V > assembledValues = valuesInKeyOrder .stream ().map (Try ::get ).collect (toList ());
357
+ return completedFuture (assembledValues );
358
358
} else {
359
359
//
360
360
// we missed some of the keys from cache, so send them to the batch loader
361
361
// and then fill in their values
362
362
//
363
- CompletableFuture <List <V >> batchLoad = invokeLoader (cacheMissedKeys , cacheMissedContexts );
363
+ CompletableFuture <List <V >> batchLoad = invokeLoader (missedKeys , missedKeyContexts );
364
364
return batchLoad .thenCompose (missedValues -> {
365
- assertResultSize (cacheMissedKeys , missedValues );
365
+ assertResultSize (missedKeys , missedValues );
366
366
367
367
for (int i = 0 ; i < missedValues .size (); i ++) {
368
- K missedKey = cacheMissedKeys .get (i );
369
368
V v = missedValues .get (i );
370
- valuesInKeyOrder .put (missedKey , v );
369
+ Integer listIndex = missedKeyIndexes .get (i );
370
+ valuesInKeyOrder .set (listIndex , Try .succeeded (v ));
371
371
}
372
- List <V > assembledValues = new ArrayList <>( valuesInKeyOrder .values ());
372
+ List <V > assembledValues = valuesInKeyOrder .stream (). map ( Try :: get ). collect ( toList ());
373
373
//
374
374
// fire off a call to the ValueCache to allow it to set values into the
375
375
// cache now that we have them
376
- return setToValueCache (assembledValues , cacheMissedKeys , missedValues );
376
+ return setToValueCache (assembledValues , missedKeys , missedValues );
377
377
});
378
378
}
379
379
});
@@ -405,7 +405,7 @@ private CompletableFuture<List<V>> invokeListBatchLoader(List<K> keys, BatchLoad
405
405
} else {
406
406
loadResult = ((BatchLoader <K , V >) batchLoadFunction ).load (keys );
407
407
}
408
- return nonNull (loadResult , () -> "Your batch loader function MUST return a non null CompletionStage promise " ).toCompletableFuture ();
408
+ return nonNull (loadResult , () -> "Your batch loader function MUST return a non null CompletionStage" ).toCompletableFuture ();
409
409
}
410
410
411
411
@@ -422,7 +422,7 @@ private CompletableFuture<List<V>> invokeMapBatchLoader(List<K> keys, BatchLoade
422
422
} else {
423
423
loadResult = ((MappedBatchLoader <K , V >) batchLoadFunction ).load (setOfKeys );
424
424
}
425
- CompletableFuture <Map <K , V >> mapBatchLoad = nonNull (loadResult , () -> "Your batch loader function MUST return a non null CompletionStage promise " ).toCompletableFuture ();
425
+ CompletableFuture <Map <K , V >> mapBatchLoad = nonNull (loadResult , () -> "Your batch loader function MUST return a non null CompletionStage" ).toCompletableFuture ();
426
426
return mapBatchLoad .thenApply (map -> {
427
427
List <V > values = new ArrayList <>();
428
428
for (K key : keys ) {
@@ -445,7 +445,7 @@ int dispatchDepth() {
445
445
446
446
private CompletableFuture <List <Try <V >>> getFromValueCache (List <K > keys ) {
447
447
try {
448
- return nonNull (valueCache .getValues (keys ), () -> "Your ValueCache.getValues function MUST return a non null promise " );
448
+ return nonNull (valueCache .getValues (keys ), () -> "Your ValueCache.getValues function MUST return a non null CompletableFuture " );
449
449
} catch (RuntimeException e ) {
450
450
return CompletableFutureKit .failedFuture (e );
451
451
}
@@ -456,7 +456,7 @@ private CompletableFuture<List<V>> setToValueCache(List<V> assembledValues, List
456
456
boolean completeValueAfterCacheSet = loaderOptions .getValueCacheOptions ().isCompleteValueAfterCacheSet ();
457
457
if (completeValueAfterCacheSet ) {
458
458
return nonNull (valueCache
459
- .setValues (missedKeys , missedValues ), () -> "Your ValueCache.setValues function MUST return a non null promise " )
459
+ .setValues (missedKeys , missedValues ), () -> "Your ValueCache.setValues function MUST return a non null CompletableFuture " )
460
460
// we dont trust the set cache to give us the values back - we have them - lets use them
461
461
// if the cache set fails - then they wont be in cache and maybe next time they will
462
462
.handle ((ignored , setExIgnored ) -> assembledValues );
0 commit comments