diff --git a/src/main/java/org/dataloader/DataLoaderHelper.java b/src/main/java/org/dataloader/DataLoaderHelper.java index a788889..dbd9383 100644 --- a/src/main/java/org/dataloader/DataLoaderHelper.java +++ b/src/main/java/org/dataloader/DataLoaderHelper.java @@ -8,7 +8,6 @@ import java.time.Instant; import java.util.ArrayList; import java.util.Collection; -import java.util.LinkedHashMap; import java.util.LinkedHashSet; import java.util.List; import java.util.Map; @@ -335,45 +334,46 @@ CompletableFuture> invokeLoader(List keys, List keyContexts, assertState(keys.size() == cachedValues.size(), () -> "The size of the cached values MUST be the same size as the key list"); - LinkedHashMap valuesInKeyOrder = new LinkedHashMap<>(); - List cacheMissedKeys = new ArrayList<>(); - List cacheMissedContexts = new ArrayList<>(); + // the following is NOT a Map because keys in data loader can repeat (by design) + // and hence "a","b","c","b" is a valid set of keys + List> valuesInKeyOrder = new ArrayList<>(); + List missedKeyIndexes = new ArrayList<>(); + List missedKeys = new ArrayList<>(); + List missedKeyContexts = new ArrayList<>(); for (int i = 0; i < keys.size(); i++) { - K key = keys.get(i); - Object keyContext = keyContexts.get(i); Try cacheGet = cachedValues.get(i); - if (cacheGet.isSuccess()) { - valuesInKeyOrder.put(key, cacheGet.get()); - } else { - valuesInKeyOrder.put(key, null); // an entry to be replaced later - cacheMissedKeys.add(key); - cacheMissedContexts.add(keyContext); + valuesInKeyOrder.add(cacheGet); + if (cacheGet.isFailure()) { + missedKeyIndexes.add(i); + missedKeys.add(keys.get(i)); + missedKeyContexts.add(keyContexts.get(i)); } } - if (cacheMissedKeys.isEmpty()) { + if (missedKeys.isEmpty()) { // // everything was cached // - return completedFuture(new ArrayList<>(valuesInKeyOrder.values())); + List assembledValues = valuesInKeyOrder.stream().map(Try::get).collect(toList()); + return completedFuture(assembledValues); } else { // // we missed some of the keys from cache, so send them to the batch loader // and then fill in their values // - CompletableFuture> batchLoad = invokeLoader(cacheMissedKeys, cacheMissedContexts); + CompletableFuture> batchLoad = invokeLoader(missedKeys, missedKeyContexts); return batchLoad.thenCompose(missedValues -> { - assertResultSize(cacheMissedKeys, missedValues); + assertResultSize(missedKeys, missedValues); for (int i = 0; i < missedValues.size(); i++) { - K missedKey = cacheMissedKeys.get(i); V v = missedValues.get(i); - valuesInKeyOrder.put(missedKey, v); + Integer listIndex = missedKeyIndexes.get(i); + valuesInKeyOrder.set(listIndex, Try.succeeded(v)); } - List assembledValues = new ArrayList<>(valuesInKeyOrder.values()); + List assembledValues = valuesInKeyOrder.stream().map(Try::get).collect(toList()); // // fire off a call to the ValueCache to allow it to set values into the // cache now that we have them - return setToValueCache(assembledValues, cacheMissedKeys, missedValues); + return setToValueCache(assembledValues, missedKeys, missedValues); }); } }); @@ -405,7 +405,7 @@ private CompletableFuture> invokeListBatchLoader(List keys, BatchLoad } else { loadResult = ((BatchLoader) batchLoadFunction).load(keys); } - return nonNull(loadResult, () -> "Your batch loader function MUST return a non null CompletionStage promise").toCompletableFuture(); + return nonNull(loadResult, () -> "Your batch loader function MUST return a non null CompletionStage").toCompletableFuture(); } @@ -422,7 +422,7 @@ private CompletableFuture> invokeMapBatchLoader(List keys, BatchLoade } else { loadResult = ((MappedBatchLoader) batchLoadFunction).load(setOfKeys); } - CompletableFuture> mapBatchLoad = nonNull(loadResult, () -> "Your batch loader function MUST return a non null CompletionStage promise").toCompletableFuture(); + CompletableFuture> mapBatchLoad = nonNull(loadResult, () -> "Your batch loader function MUST return a non null CompletionStage").toCompletableFuture(); return mapBatchLoad.thenApply(map -> { List values = new ArrayList<>(); for (K key : keys) { @@ -445,7 +445,7 @@ int dispatchDepth() { private CompletableFuture>> getFromValueCache(List keys) { try { - return nonNull(valueCache.getValues(keys), () -> "Your ValueCache.getValues function MUST return a non null promise"); + return nonNull(valueCache.getValues(keys), () -> "Your ValueCache.getValues function MUST return a non null CompletableFuture"); } catch (RuntimeException e) { return CompletableFutureKit.failedFuture(e); } @@ -456,7 +456,7 @@ private CompletableFuture> setToValueCache(List assembledValues, List boolean completeValueAfterCacheSet = loaderOptions.getValueCacheOptions().isCompleteValueAfterCacheSet(); if (completeValueAfterCacheSet) { return nonNull(valueCache - .setValues(missedKeys, missedValues), () -> "Your ValueCache.setValues function MUST return a non null promise") + .setValues(missedKeys, missedValues), () -> "Your ValueCache.setValues function MUST return a non null CompletableFuture") // we dont trust the set cache to give us the values back - we have them - lets use them // if the cache set fails - then they wont be in cache and maybe next time they will .handle((ignored, setExIgnored) -> assembledValues);