@@ -224,9 +224,8 @@ private CompletableFuture<List<V>> sliceIntoBatchesOfBatches(List<K> keys, List<
224
224
@ SuppressWarnings ("unchecked" )
225
225
private CompletableFuture <List <V >> dispatchQueueBatch (List <K > keys , List <Object > callContexts , List <CompletableFuture <V >> queuedFutures ) {
226
226
stats .incrementBatchLoadCountBy (keys .size ());
227
- CompletionStage <List <V >> batchLoad = invokeLoader (keys , callContexts , loaderOptions .cachingEnabled ());
227
+ CompletableFuture <List <V >> batchLoad = invokeLoader (keys , callContexts , loaderOptions .cachingEnabled ());
228
228
return batchLoad
229
- .toCompletableFuture ()
230
229
.thenApply (values -> {
231
230
assertResultSize (keys , values );
232
231
@@ -327,7 +326,7 @@ CompletableFuture<V> invokeLoaderImmediately(K key, Object keyContext, boolean c
327
326
.toCompletableFuture ();
328
327
}
329
328
330
- CompletionStage <List <V >> invokeLoader (List <K > keys , List <Object > keyContexts , boolean cachingEnabled ) {
329
+ CompletableFuture <List <V >> invokeLoader (List <K > keys , List <Object > keyContexts , boolean cachingEnabled ) {
331
330
if (!cachingEnabled ) {
332
331
return invokeLoader (keys , keyContexts );
333
332
}
@@ -361,7 +360,7 @@ CompletionStage<List<V>> invokeLoader(List<K> keys, List<Object> keyContexts, bo
361
360
// we missed some of the keys from cache, so send them to the batch loader
362
361
// and then fill in their values
363
362
//
364
- CompletionStage <List <V >> batchLoad = invokeLoader (cacheMissedKeys , cacheMissedContexts );
363
+ CompletableFuture <List <V >> batchLoad = invokeLoader (cacheMissedKeys , cacheMissedContexts );
365
364
return batchLoad .thenCompose (missedValues -> {
366
365
assertResultSize (cacheMissedKeys , missedValues );
367
366
@@ -381,8 +380,8 @@ CompletionStage<List<V>> invokeLoader(List<K> keys, List<Object> keyContexts, bo
381
380
}
382
381
383
382
384
- CompletionStage <List <V >> invokeLoader (List <K > keys , List <Object > keyContexts ) {
385
- CompletionStage <List <V >> batchLoad ;
383
+ CompletableFuture <List <V >> invokeLoader (List <K > keys , List <Object > keyContexts ) {
384
+ CompletableFuture <List <V >> batchLoad ;
386
385
try {
387
386
Object context = loaderOptions .getBatchLoaderContextProvider ().getContext ();
388
387
BatchLoaderEnvironment environment = BatchLoaderEnvironment .newBatchLoaderEnvironment ()
@@ -399,14 +398,14 @@ CompletionStage<List<V>> invokeLoader(List<K> keys, List<Object> keyContexts) {
399
398
}
400
399
401
400
@ SuppressWarnings ("unchecked" )
402
- private CompletionStage <List <V >> invokeListBatchLoader (List <K > keys , BatchLoaderEnvironment environment ) {
401
+ private CompletableFuture <List <V >> invokeListBatchLoader (List <K > keys , BatchLoaderEnvironment environment ) {
403
402
CompletionStage <List <V >> loadResult ;
404
403
if (batchLoadFunction instanceof BatchLoaderWithContext ) {
405
404
loadResult = ((BatchLoaderWithContext <K , V >) batchLoadFunction ).load (keys , environment );
406
405
} else {
407
406
loadResult = ((BatchLoader <K , V >) batchLoadFunction ).load (keys );
408
407
}
409
- return nonNull (loadResult , () -> "Your batch loader function MUST return a non null CompletionStage promise" );
408
+ return nonNull (loadResult , () -> "Your batch loader function MUST return a non null CompletionStage promise" ). toCompletableFuture () ;
410
409
}
411
410
412
411
@@ -415,15 +414,15 @@ private CompletionStage<List<V>> invokeListBatchLoader(List<K> keys, BatchLoader
415
414
* to missing elements.
416
415
*/
417
416
@ SuppressWarnings ("unchecked" )
418
- private CompletionStage <List <V >> invokeMapBatchLoader (List <K > keys , BatchLoaderEnvironment environment ) {
417
+ private CompletableFuture <List <V >> invokeMapBatchLoader (List <K > keys , BatchLoaderEnvironment environment ) {
419
418
CompletionStage <Map <K , V >> loadResult ;
420
419
Set <K > setOfKeys = new LinkedHashSet <>(keys );
421
420
if (batchLoadFunction instanceof MappedBatchLoaderWithContext ) {
422
421
loadResult = ((MappedBatchLoaderWithContext <K , V >) batchLoadFunction ).load (setOfKeys , environment );
423
422
} else {
424
423
loadResult = ((MappedBatchLoader <K , V >) batchLoadFunction ).load (setOfKeys );
425
424
}
426
- CompletionStage <Map <K , V >> mapBatchLoad = nonNull (loadResult , () -> "Your batch loader function MUST return a non null CompletionStage promise" );
425
+ CompletableFuture <Map <K , V >> mapBatchLoad = nonNull (loadResult , () -> "Your batch loader function MUST return a non null CompletionStage promise" ). toCompletableFuture ( );
427
426
return mapBatchLoad .thenApply (map -> {
428
427
List <V > values = new ArrayList <>();
429
428
for (K key : keys ) {
@@ -452,7 +451,7 @@ private CompletableFuture<List<Try<V>>> getFromValueCache(List<K> keys) {
452
451
}
453
452
}
454
453
455
- private CompletionStage <List <V >> setToValueCache (List <V > assembledValues , List <K > missedKeys , List <V > missedValues ) {
454
+ private CompletableFuture <List <V >> setToValueCache (List <V > assembledValues , List <K > missedKeys , List <V > missedValues ) {
456
455
try {
457
456
boolean completeValueAfterCacheSet = loaderOptions .getValueCacheOptions ().isCompleteValueAfterCacheSet ();
458
457
if (completeValueAfterCacheSet ) {
0 commit comments