diff --git a/build.gradle b/build.gradle index 6dea10e..edc81c7 100644 --- a/build.gradle +++ b/build.gradle @@ -30,7 +30,6 @@ def getDevelopmentVersion() { version } - def releaseVersion = System.env.RELEASE_VERSION version = releaseVersion ? releaseVersion : getDevelopmentVersion() group = 'com.graphql-java' @@ -76,6 +75,7 @@ dependencies { testImplementation 'org.awaitility:awaitility:2.0.0' testImplementation 'io.projectreactor:reactor-core:3.6.6' testImplementation 'com.github.ben-manes.caffeine:caffeine:2.9.0' + testImplementation 'io.projectreactor:reactor-core:3.6.6' } task sourcesJar(type: Jar) { diff --git a/src/main/java/org/dataloader/BatchPublisher.java b/src/main/java/org/dataloader/BatchPublisher.java new file mode 100644 index 0000000..5ab41e1 --- /dev/null +++ b/src/main/java/org/dataloader/BatchPublisher.java @@ -0,0 +1,22 @@ +package org.dataloader; + +import org.reactivestreams.Subscriber; + +import java.util.List; + +/** + * A function that is invoked for batch loading a stream of data values indicated by the provided list of keys. + *

+ * The function will call the provided {@link Subscriber} to process the values it has retrieved to allow + * the future returned by {@link DataLoader#load(Object)} to complete as soon as the individual value is available + * (rather than when all values have been retrieved). + *

+ * NOTE: It is required that {@link Subscriber#onNext(V)} is invoked on each value in the same order as + * the provided keys. + * + * @param type parameter indicating the type of keys to use for data load requests. + * @param type parameter indicating the type of values returned + */ +public interface BatchPublisher { + void load(List keys, Subscriber subscriber); +} diff --git a/src/main/java/org/dataloader/BatchPublisherWithContext.java b/src/main/java/org/dataloader/BatchPublisherWithContext.java new file mode 100644 index 0000000..effda90 --- /dev/null +++ b/src/main/java/org/dataloader/BatchPublisherWithContext.java @@ -0,0 +1,12 @@ +package org.dataloader; + +import org.reactivestreams.Subscriber; + +import java.util.List; + +/** + * An {@link BatchPublisher} with a {@link BatchLoaderEnvironment} provided as an extra parameter to {@link #load}. + */ +public interface BatchPublisherWithContext { + void load(List keys, Subscriber subscriber, BatchLoaderEnvironment environment); +} diff --git a/src/main/java/org/dataloader/DataLoaderFactory.java b/src/main/java/org/dataloader/DataLoaderFactory.java index 013f473..db14f2e 100644 --- a/src/main/java/org/dataloader/DataLoaderFactory.java +++ b/src/main/java/org/dataloader/DataLoaderFactory.java @@ -278,6 +278,274 @@ public static DataLoader newMappedDataLoaderWithTry(MappedBatchLoad return mkDataLoader(batchLoadFunction, options); } + /** + * Creates new DataLoader with the specified batch loader function and default options + * (batching, caching and unlimited batch size). + * + * @param batchLoadFunction the batch load function to use + * @param the key type + * @param the value type + * + * @return a new DataLoader + */ + public static DataLoader newPublisherDataLoader(BatchPublisher batchLoadFunction) { + return newPublisherDataLoader(batchLoadFunction, null); + } + + /** + * Creates new DataLoader with the specified batch loader function with the provided options + * + * @param batchLoadFunction the batch load function to use + * @param options the options to use + * @param the key type + * @param the value type + * + * @return a new DataLoader + */ + public static DataLoader newPublisherDataLoader(BatchPublisher batchLoadFunction, DataLoaderOptions options) { + return mkDataLoader(batchLoadFunction, options); + } + + /** + * Creates new DataLoader with the specified batch loader function and default options + * (batching, caching and unlimited batch size) where the batch loader function returns a list of + * {@link org.dataloader.Try} objects. + *

+ * If it's important you to know the exact status of each item in a batch call and whether it threw exceptions then + * you can use this form to create the data loader. + *

+ * Using Try objects allows you to capture a value returned or an exception that might + * have occurred trying to get a value. . + * + * @param batchLoadFunction the batch load function to use that uses {@link org.dataloader.Try} objects + * @param the key type + * @param the value type + * + * @return a new DataLoader + */ + public static DataLoader newPublisherDataLoaderWithTry(BatchPublisher> batchLoadFunction) { + return newPublisherDataLoaderWithTry(batchLoadFunction, null); + } + + /** + * Creates new DataLoader with the specified batch loader function and with the provided options + * where the batch loader function returns a list of + * {@link org.dataloader.Try} objects. + * + * @param batchLoadFunction the batch load function to use that uses {@link org.dataloader.Try} objects + * @param options the options to use + * @param the key type + * @param the value type + * + * @return a new DataLoader + * + * @see #newDataLoaderWithTry(BatchLoader) + */ + public static DataLoader newPublisherDataLoaderWithTry(BatchPublisher> batchLoadFunction, DataLoaderOptions options) { + return mkDataLoader(batchLoadFunction, options); + } + + /** + * Creates new DataLoader with the specified batch loader function and default options + * (batching, caching and unlimited batch size). + * + * @param batchLoadFunction the batch load function to use + * @param the key type + * @param the value type + * + * @return a new DataLoader + */ + public static DataLoader newPublisherDataLoader(BatchPublisherWithContext batchLoadFunction) { + return newPublisherDataLoader(batchLoadFunction, null); + } + + /** + * Creates new DataLoader with the specified batch loader function with the provided options + * + * @param batchLoadFunction the batch load function to use + * @param options the options to use + * @param the key type + * @param the value type + * + * @return a new DataLoader + */ + public static DataLoader newPublisherDataLoader(BatchPublisherWithContext batchLoadFunction, DataLoaderOptions options) { + return mkDataLoader(batchLoadFunction, options); + } + + /** + * Creates new DataLoader with the specified batch loader function and default options + * (batching, caching and unlimited batch size) where the batch loader function returns a list of + * {@link org.dataloader.Try} objects. + *

+ * If it's important you to know the exact status of each item in a batch call and whether it threw exceptions then + * you can use this form to create the data loader. + *

+ * Using Try objects allows you to capture a value returned or an exception that might + * have occurred trying to get a value. . + * + * @param batchLoadFunction the batch load function to use that uses {@link org.dataloader.Try} objects + * @param the key type + * @param the value type + * + * @return a new DataLoader + */ + public static DataLoader newPublisherDataLoaderWithTry(BatchPublisherWithContext> batchLoadFunction) { + return newPublisherDataLoaderWithTry(batchLoadFunction, null); + } + + /** + * Creates new DataLoader with the specified batch loader function and with the provided options + * where the batch loader function returns a list of + * {@link org.dataloader.Try} objects. + * + * @param batchLoadFunction the batch load function to use that uses {@link org.dataloader.Try} objects + * @param options the options to use + * @param the key type + * @param the value type + * + * @return a new DataLoader + * + * @see #newPublisherDataLoaderWithTry(BatchPublisher) + */ + public static DataLoader newPublisherDataLoaderWithTry(BatchPublisherWithContext> batchLoadFunction, DataLoaderOptions options) { + return mkDataLoader(batchLoadFunction, options); + } + + /** + * Creates new DataLoader with the specified batch loader function and default options + * (batching, caching and unlimited batch size). + * + * @param batchLoadFunction the batch load function to use + * @param the key type + * @param the value type + * + * @return a new DataLoader + */ + public static DataLoader newMappedPublisherDataLoader(MappedBatchPublisher batchLoadFunction) { + return newMappedPublisherDataLoader(batchLoadFunction, null); + } + + /** + * Creates new DataLoader with the specified batch loader function with the provided options + * + * @param batchLoadFunction the batch load function to use + * @param options the options to use + * @param the key type + * @param the value type + * + * @return a new DataLoader + */ + public static DataLoader newMappedPublisherDataLoader(MappedBatchPublisher batchLoadFunction, DataLoaderOptions options) { + return mkDataLoader(batchLoadFunction, options); + } + + /** + * Creates new DataLoader with the specified batch loader function and default options + * (batching, caching and unlimited batch size) where the batch loader function returns a list of + * {@link org.dataloader.Try} objects. + *

+ * If it's important you to know the exact status of each item in a batch call and whether it threw exceptions then + * you can use this form to create the data loader. + *

+ * Using Try objects allows you to capture a value returned or an exception that might + * have occurred trying to get a value. . + * + * @param batchLoadFunction the batch load function to use that uses {@link org.dataloader.Try} objects + * @param the key type + * @param the value type + * + * @return a new DataLoader + */ + public static DataLoader newMappedPublisherDataLoaderWithTry(MappedBatchPublisher> batchLoadFunction) { + return newMappedPublisherDataLoaderWithTry(batchLoadFunction, null); + } + + /** + * Creates new DataLoader with the specified batch loader function and with the provided options + * where the batch loader function returns a list of + * {@link org.dataloader.Try} objects. + * + * @param batchLoadFunction the batch load function to use that uses {@link org.dataloader.Try} objects + * @param options the options to use + * @param the key type + * @param the value type + * + * @return a new DataLoader + * + * @see #newDataLoaderWithTry(BatchLoader) + */ + public static DataLoader newMappedPublisherDataLoaderWithTry(MappedBatchPublisher> batchLoadFunction, DataLoaderOptions options) { + return mkDataLoader(batchLoadFunction, options); + } + + /** + * Creates new DataLoader with the specified batch loader function and default options + * (batching, caching and unlimited batch size). + * + * @param batchLoadFunction the batch load function to use + * @param the key type + * @param the value type + * + * @return a new DataLoader + */ + public static DataLoader newMappedPublisherDataLoader(MappedBatchPublisherWithContext batchLoadFunction) { + return newMappedPublisherDataLoader(batchLoadFunction, null); + } + + /** + * Creates new DataLoader with the specified batch loader function with the provided options + * + * @param batchLoadFunction the batch load function to use + * @param options the options to use + * @param the key type + * @param the value type + * + * @return a new DataLoader + */ + public static DataLoader newMappedPublisherDataLoader(MappedBatchPublisherWithContext batchLoadFunction, DataLoaderOptions options) { + return mkDataLoader(batchLoadFunction, options); + } + + /** + * Creates new DataLoader with the specified batch loader function and default options + * (batching, caching and unlimited batch size) where the batch loader function returns a list of + * {@link org.dataloader.Try} objects. + *

+ * If it's important you to know the exact status of each item in a batch call and whether it threw exceptions then + * you can use this form to create the data loader. + *

+ * Using Try objects allows you to capture a value returned or an exception that might + * have occurred trying to get a value. . + * + * @param batchLoadFunction the batch load function to use that uses {@link org.dataloader.Try} objects + * @param the key type + * @param the value type + * + * @return a new DataLoader + */ + public static DataLoader newMappedPublisherDataLoaderWithTry(MappedBatchPublisherWithContext> batchLoadFunction) { + return newMappedPublisherDataLoaderWithTry(batchLoadFunction, null); + } + + /** + * Creates new DataLoader with the specified batch loader function and with the provided options + * where the batch loader function returns a list of + * {@link org.dataloader.Try} objects. + * + * @param batchLoadFunction the batch load function to use that uses {@link org.dataloader.Try} objects + * @param options the options to use + * @param the key type + * @param the value type + * + * @return a new DataLoader + * + * @see #newMappedPublisherDataLoaderWithTry(MappedBatchPublisher) + */ + public static DataLoader newMappedPublisherDataLoaderWithTry(MappedBatchPublisherWithContext> batchLoadFunction, DataLoaderOptions options) { + return mkDataLoader(batchLoadFunction, options); + } + static DataLoader mkDataLoader(Object batchLoadFunction, DataLoaderOptions options) { return new DataLoader<>(batchLoadFunction, options); } diff --git a/src/main/java/org/dataloader/DataLoaderHelper.java b/src/main/java/org/dataloader/DataLoaderHelper.java index d934de2..ee8d78b 100644 --- a/src/main/java/org/dataloader/DataLoaderHelper.java +++ b/src/main/java/org/dataloader/DataLoaderHelper.java @@ -10,11 +10,14 @@ import org.dataloader.stats.context.IncrementCacheHitCountStatisticsContext; import org.dataloader.stats.context.IncrementLoadCountStatisticsContext; import org.dataloader.stats.context.IncrementLoadErrorCountStatisticsContext; +import org.reactivestreams.Subscriber; +import org.reactivestreams.Subscription; import java.time.Clock; import java.time.Instant; import java.util.ArrayList; import java.util.Collection; +import java.util.HashMap; import java.util.LinkedHashSet; import java.util.List; import java.util.Map; @@ -241,10 +244,14 @@ private CompletableFuture> sliceIntoBatchesOfBatches(List keys, List< @SuppressWarnings("unchecked") private CompletableFuture> dispatchQueueBatch(List keys, List callContexts, List> queuedFutures) { stats.incrementBatchLoadCountBy(keys.size(), new IncrementBatchLoadCountByStatisticsContext<>(keys, callContexts)); - CompletableFuture> batchLoad = invokeLoader(keys, callContexts, loaderOptions.cachingEnabled()); + CompletableFuture> batchLoad = invokeLoader(keys, callContexts, queuedFutures, loaderOptions.cachingEnabled()); return batchLoad .thenApply(values -> { assertResultSize(keys, values); + if (isPublisher() || isMappedPublisher()) { + // We have already completed the queued futures by the time the overall batchLoad future has completed. + return values; + } List clearCacheKeys = new ArrayList<>(); for (int idx = 0; idx < queuedFutures.size(); idx++) { @@ -342,14 +349,15 @@ private CompletableFuture queueOrInvokeLoader(K key, Object loadContext, bool CompletableFuture invokeLoaderImmediately(K key, Object keyContext, boolean cachingEnabled) { List keys = singletonList(key); List keyContexts = singletonList(keyContext); - return invokeLoader(keys, keyContexts, cachingEnabled) + List> queuedFutures = singletonList(new CompletableFuture<>()); + return invokeLoader(keys, keyContexts, queuedFutures, cachingEnabled) .thenApply(list -> list.get(0)) .toCompletableFuture(); } - CompletableFuture> invokeLoader(List keys, List keyContexts, boolean cachingEnabled) { + CompletableFuture> invokeLoader(List keys, List keyContexts, List> queuedFutures, boolean cachingEnabled) { if (!cachingEnabled) { - return invokeLoader(keys, keyContexts); + return invokeLoader(keys, keyContexts, queuedFutures); } CompletableFuture>> cacheCallCF = getFromValueCache(keys); return cacheCallCF.thenCompose(cachedValues -> { @@ -360,6 +368,7 @@ CompletableFuture> invokeLoader(List keys, List keyContexts, List missedKeyIndexes = new ArrayList<>(); List missedKeys = new ArrayList<>(); List missedKeyContexts = new ArrayList<>(); + List> missedQueuedFutures = new ArrayList<>(); // if they return a ValueCachingNotSupported exception then we insert this special marker value, and it // means it's a total miss, we need to get all these keys via the batch loader @@ -369,6 +378,7 @@ CompletableFuture> invokeLoader(List keys, List keyContexts, missedKeyIndexes.add(i); missedKeys.add(keys.get(i)); missedKeyContexts.add(keyContexts.get(i)); + missedQueuedFutures.add(queuedFutures.get(i)); } } else { assertState(keys.size() == cachedValues.size(), () -> "The size of the cached values MUST be the same size as the key list"); @@ -393,7 +403,7 @@ CompletableFuture> invokeLoader(List keys, List keyContexts, // we missed some keys from cache, so send them to the batch loader // and then fill in their values // - CompletableFuture> batchLoad = invokeLoader(missedKeys, missedKeyContexts); + CompletableFuture> batchLoad = invokeLoader(missedKeys, missedKeyContexts, missedQueuedFutures); return batchLoad.thenCompose(missedValues -> { assertResultSize(missedKeys, missedValues); @@ -412,8 +422,7 @@ CompletableFuture> invokeLoader(List keys, List keyContexts, }); } - - CompletableFuture> invokeLoader(List keys, List keyContexts) { + CompletableFuture> invokeLoader(List keys, List keyContexts, List> queuedFutures) { CompletableFuture> batchLoad; try { Object context = loaderOptions.getBatchLoaderContextProvider().getContext(); @@ -421,6 +430,10 @@ CompletableFuture> invokeLoader(List keys, List keyContexts) .context(context).keyContexts(keys, keyContexts).build(); if (isMapLoader()) { batchLoad = invokeMapBatchLoader(keys, environment); + } else if (isPublisher()) { + batchLoad = invokeBatchPublisher(keys, keyContexts, queuedFutures, environment); + } else if (isMappedPublisher()) { + batchLoad = invokeMappedBatchPublisher(keys, keyContexts, queuedFutures, environment); } else { batchLoad = invokeListBatchLoader(keys, environment); } @@ -492,10 +505,68 @@ private CompletableFuture> invokeMapBatchLoader(List keys, BatchLoade }); } + private CompletableFuture> invokeBatchPublisher(List keys, List keyContexts, List> queuedFutures, BatchLoaderEnvironment environment) { + CompletableFuture> loadResult = new CompletableFuture<>(); + Subscriber subscriber = new DataLoaderSubscriber(loadResult, keys, keyContexts, queuedFutures); + + BatchLoaderScheduler batchLoaderScheduler = loaderOptions.getBatchLoaderScheduler(); + if (batchLoadFunction instanceof BatchPublisherWithContext) { + BatchPublisherWithContext loadFunction = (BatchPublisherWithContext) batchLoadFunction; + if (batchLoaderScheduler != null) { + BatchLoaderScheduler.ScheduledBatchPublisherCall loadCall = () -> loadFunction.load(keys, subscriber, environment); + batchLoaderScheduler.scheduleBatchPublisher(loadCall, keys, environment); + } else { + loadFunction.load(keys, subscriber, environment); + } + } else { + BatchPublisher loadFunction = (BatchPublisher) batchLoadFunction; + if (batchLoaderScheduler != null) { + BatchLoaderScheduler.ScheduledBatchPublisherCall loadCall = () -> loadFunction.load(keys, subscriber); + batchLoaderScheduler.scheduleBatchPublisher(loadCall, keys, null); + } else { + loadFunction.load(keys, subscriber); + } + } + return loadResult; + } + + private CompletableFuture> invokeMappedBatchPublisher(List keys, List keyContexts, List> queuedFutures, BatchLoaderEnvironment environment) { + CompletableFuture> loadResult = new CompletableFuture<>(); + Subscriber> subscriber = new DataLoaderMapEntrySubscriber(loadResult, keys, keyContexts, queuedFutures); + + BatchLoaderScheduler batchLoaderScheduler = loaderOptions.getBatchLoaderScheduler(); + if (batchLoadFunction instanceof MappedBatchPublisherWithContext) { + MappedBatchPublisherWithContext loadFunction = (MappedBatchPublisherWithContext) batchLoadFunction; + if (batchLoaderScheduler != null) { + BatchLoaderScheduler.ScheduledBatchPublisherCall loadCall = () -> loadFunction.load(keys, subscriber, environment); + batchLoaderScheduler.scheduleBatchPublisher(loadCall, keys, environment); + } else { + loadFunction.load(keys, subscriber, environment); + } + } else { + MappedBatchPublisher loadFunction = (MappedBatchPublisher) batchLoadFunction; + if (batchLoaderScheduler != null) { + BatchLoaderScheduler.ScheduledBatchPublisherCall loadCall = () -> loadFunction.load(keys, subscriber); + batchLoaderScheduler.scheduleBatchPublisher(loadCall, keys, null); + } else { + loadFunction.load(keys, subscriber); + } + } + return loadResult; + } + private boolean isMapLoader() { return batchLoadFunction instanceof MappedBatchLoader || batchLoadFunction instanceof MappedBatchLoaderWithContext; } + private boolean isPublisher() { + return batchLoadFunction instanceof BatchPublisher; + } + + private boolean isMappedPublisher() { + return batchLoadFunction instanceof MappedBatchPublisher; + } + int dispatchDepth() { synchronized (dataLoader) { return loaderQueue.size(); @@ -546,4 +617,197 @@ private CompletableFuture> setToValueCache(List assembledValues, List private static DispatchResult emptyDispatchResult() { return (DispatchResult) EMPTY_DISPATCH_RESULT; } + + private class DataLoaderSubscriber implements Subscriber { + + private final CompletableFuture> valuesFuture; + private final List keys; + private final List callContexts; + private final List> queuedFutures; + + private final List clearCacheKeys = new ArrayList<>(); + private final List completedValues = new ArrayList<>(); + private int idx = 0; + private boolean onErrorCalled = false; + private boolean onCompleteCalled = false; + + private DataLoaderSubscriber( + CompletableFuture> valuesFuture, + List keys, + List callContexts, + List> queuedFutures + ) { + this.valuesFuture = valuesFuture; + this.keys = keys; + this.callContexts = callContexts; + this.queuedFutures = queuedFutures; + } + + @Override + public void onSubscribe(Subscription subscription) { + subscription.request(keys.size()); + } + + // onNext may be called by multiple threads - for the time being, we pass 'synchronized' to guarantee + // correctness (at the cost of speed). + @Override + public synchronized void onNext(V value) { + assertState(!onErrorCalled, () -> "onError has already been called; onNext may not be invoked."); + assertState(!onCompleteCalled, () -> "onComplete has already been called; onNext may not be invoked."); + + K key = keys.get(idx); + Object callContext = callContexts.get(idx); + CompletableFuture future = queuedFutures.get(idx); + if (value instanceof Try) { + // we allow the batch loader to return a Try so we can better represent a computation + // that might have worked or not. + Try tryValue = (Try) value; + if (tryValue.isSuccess()) { + future.complete(tryValue.get()); + } else { + stats.incrementLoadErrorCount(new IncrementLoadErrorCountStatisticsContext<>(key, callContext)); + future.completeExceptionally(tryValue.getThrowable()); + clearCacheKeys.add(keys.get(idx)); + } + } else { + future.complete(value); + } + + completedValues.add(value); + idx++; + } + + @Override + public void onComplete() { + assertState(!onErrorCalled, () -> "onError has already been called; onComplete may not be invoked."); + onCompleteCalled = true; + + assertResultSize(keys, completedValues); + + possiblyClearCacheEntriesOnExceptions(clearCacheKeys); + valuesFuture.complete(completedValues); + } + + @Override + public void onError(Throwable ex) { + assertState(!onCompleteCalled, () -> "onComplete has already been called; onError may not be invoked."); + onErrorCalled = true; + + stats.incrementBatchLoadExceptionCount(new IncrementBatchLoadExceptionCountStatisticsContext<>(keys, callContexts)); + if (ex instanceof CompletionException) { + ex = ex.getCause(); + } + // Set the remaining keys to the exception. + for (int i = idx; i < queuedFutures.size(); i++) { + K key = keys.get(i); + CompletableFuture future = queuedFutures.get(i); + future.completeExceptionally(ex); + // clear any cached view of this key because they all failed + dataLoader.clear(key); + } + } + } + + private class DataLoaderMapEntrySubscriber implements Subscriber> { + private final CompletableFuture> valuesFuture; + private final List keys; + private final List callContexts; + private final List> queuedFutures; + private final Map callContextByKey; + private final Map> queuedFutureByKey; + + private final List clearCacheKeys = new ArrayList<>(); + private final Map completedValuesByKey = new HashMap<>(); + private boolean onErrorCalled = false; + private boolean onCompleteCalled = false; + + private DataLoaderMapEntrySubscriber( + CompletableFuture> valuesFuture, + List keys, + List callContexts, + List> queuedFutures + ) { + this.valuesFuture = valuesFuture; + this.keys = keys; + this.callContexts = callContexts; + this.queuedFutures = queuedFutures; + + this.callContextByKey = new HashMap<>(); + this.queuedFutureByKey = new HashMap<>(); + for (int idx = 0; idx < queuedFutures.size(); idx++) { + K key = keys.get(idx); + Object callContext = callContexts.get(idx); + CompletableFuture queuedFuture = queuedFutures.get(idx); + callContextByKey.put(key, callContext); + queuedFutureByKey.put(key, queuedFuture); + } + } + + @Override + public void onSubscribe(Subscription subscription) { + subscription.request(keys.size()); + } + + @Override + public void onNext(Map.Entry entry) { + assertState(!onErrorCalled, () -> "onError has already been called; onNext may not be invoked."); + assertState(!onCompleteCalled, () -> "onComplete has already been called; onNext may not be invoked."); + K key = entry.getKey(); + V value = entry.getValue(); + + Object callContext = callContextByKey.get(key); + CompletableFuture future = queuedFutureByKey.get(key); + if (value instanceof Try) { + // we allow the batch loader to return a Try so we can better represent a computation + // that might have worked or not. + Try tryValue = (Try) value; + if (tryValue.isSuccess()) { + future.complete(tryValue.get()); + } else { + stats.incrementLoadErrorCount(new IncrementLoadErrorCountStatisticsContext<>(key, callContext)); + future.completeExceptionally(tryValue.getThrowable()); + clearCacheKeys.add(key); + } + } else { + future.complete(value); + } + + completedValuesByKey.put(key, value); + } + + @Override + public void onComplete() { + assertState(!onErrorCalled, () -> "onError has already been called; onComplete may not be invoked."); + onCompleteCalled = true; + + possiblyClearCacheEntriesOnExceptions(clearCacheKeys); + List values = new ArrayList<>(keys.size()); + for (K key : keys) { + V value = completedValuesByKey.get(key); + values.add(value); + } + valuesFuture.complete(values); + } + + @Override + public void onError(Throwable ex) { + assertState(!onCompleteCalled, () -> "onComplete has already been called; onError may not be invoked."); + onErrorCalled = true; + + stats.incrementBatchLoadExceptionCount(new IncrementBatchLoadExceptionCountStatisticsContext<>(keys, callContexts)); + if (ex instanceof CompletionException) { + ex = ex.getCause(); + } + // Complete the futures for the remaining keys with the exception. + for (int idx = 0; idx < queuedFutures.size(); idx++) { + K key = keys.get(idx); + CompletableFuture future = queuedFutureByKey.get(key); + if (!completedValuesByKey.containsKey(key)) { + future.completeExceptionally(ex); + // clear any cached view of this key because they all failed + dataLoader.clear(key); + } + } + } + } } diff --git a/src/main/java/org/dataloader/MappedBatchPublisher.java b/src/main/java/org/dataloader/MappedBatchPublisher.java new file mode 100644 index 0000000..9b3fcb9 --- /dev/null +++ b/src/main/java/org/dataloader/MappedBatchPublisher.java @@ -0,0 +1,20 @@ +package org.dataloader; + +import org.reactivestreams.Subscriber; + +import java.util.List; +import java.util.Map; + +/** + * A function that is invoked for batch loading a stream of data values indicated by the provided list of keys. + *

+ * The function will call the provided {@link Subscriber} to process the key/value pairs it has retrieved to allow + * the future returned by {@link DataLoader#load(Object)} to complete as soon as the individual value is available + * (rather than when all values have been retrieved). + * + * @param type parameter indicating the type of keys to use for data load requests. + * @param type parameter indicating the type of values returned + */ +public interface MappedBatchPublisher { + void load(List keys, Subscriber> subscriber); +} diff --git a/src/main/java/org/dataloader/MappedBatchPublisherWithContext.java b/src/main/java/org/dataloader/MappedBatchPublisherWithContext.java new file mode 100644 index 0000000..4810111 --- /dev/null +++ b/src/main/java/org/dataloader/MappedBatchPublisherWithContext.java @@ -0,0 +1,13 @@ +package org.dataloader; + +import org.reactivestreams.Subscriber; + +import java.util.List; +import java.util.Map; + +/** + * A {@link MappedBatchPublisher} with a {@link BatchLoaderEnvironment} provided as an extra parameter to {@link #load}. + */ +public interface MappedBatchPublisherWithContext { + void load(List keys, Subscriber> subscriber, BatchLoaderEnvironment environment); +} diff --git a/src/main/java/org/dataloader/scheduler/BatchLoaderScheduler.java b/src/main/java/org/dataloader/scheduler/BatchLoaderScheduler.java index 7cddd54..e7e95d9 100644 --- a/src/main/java/org/dataloader/scheduler/BatchLoaderScheduler.java +++ b/src/main/java/org/dataloader/scheduler/BatchLoaderScheduler.java @@ -5,6 +5,8 @@ import org.dataloader.DataLoader; import org.dataloader.DataLoaderOptions; import org.dataloader.MappedBatchLoader; +import org.dataloader.MappedBatchPublisher; +import org.dataloader.BatchPublisher; import java.util.List; import java.util.Map; @@ -42,6 +44,13 @@ interface ScheduledMappedBatchLoaderCall { CompletionStage> invoke(); } + /** + * This represents a callback that will invoke a {@link BatchPublisher} or {@link MappedBatchPublisher} function under the covers + */ + interface ScheduledBatchPublisherCall { + void invoke(); + } + /** * This is called to schedule a {@link BatchLoader} call. * @@ -71,4 +80,16 @@ interface ScheduledMappedBatchLoaderCall { * @return a promise to the values that come from the {@link BatchLoader} */ CompletionStage> scheduleMappedBatchLoader(ScheduledMappedBatchLoaderCall scheduledCall, List keys, BatchLoaderEnvironment environment); + + /** + * This is called to schedule a {@link BatchPublisher} call. + * + * @param scheduledCall the callback that needs to be invoked to allow the {@link BatchPublisher} to proceed. + * @param keys this is the list of keys that will be passed to the {@link BatchPublisher}. + * This is provided only for informative reasons and, you can't change the keys that are used + * @param environment this is the {@link BatchLoaderEnvironment} in place, + * which can be null if it's a simple {@link BatchPublisher} call + * @param the key type + */ + void scheduleBatchPublisher(ScheduledBatchPublisherCall scheduledCall, List keys, BatchLoaderEnvironment environment); } diff --git a/src/test/java/ReadmeExamples.java b/src/test/java/ReadmeExamples.java index d25dfa7..31354ea 100644 --- a/src/test/java/ReadmeExamples.java +++ b/src/test/java/ReadmeExamples.java @@ -304,6 +304,12 @@ public CompletionStage> scheduleMappedBatchLoader(ScheduledMapp return scheduledCall.invoke(); }).thenCompose(Function.identity()); } + + @Override + public void scheduleBatchPublisher(ScheduledBatchPublisherCall scheduledCall, List keys, BatchLoaderEnvironment environment) { + snooze(10); + scheduledCall.invoke(); + } }; } diff --git a/src/test/java/org/dataloader/DataLoaderBatchPublisherTest.java b/src/test/java/org/dataloader/DataLoaderBatchPublisherTest.java new file mode 100644 index 0000000..84a8b18 --- /dev/null +++ b/src/test/java/org/dataloader/DataLoaderBatchPublisherTest.java @@ -0,0 +1,1096 @@ +package org.dataloader; + +import org.dataloader.fixtures.CustomCacheMap; +import org.dataloader.fixtures.JsonObject; +import org.dataloader.fixtures.User; +import org.dataloader.fixtures.UserManager; +import org.dataloader.impl.CompletableFutureKit; +import org.junit.Test; +import reactor.core.publisher.Flux; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.List; +import java.util.Optional; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionStage; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.function.Supplier; +import java.util.stream.Collectors; +import java.util.stream.Stream; + +import static java.util.Arrays.asList; +import static java.util.Collections.emptyList; +import static java.util.Collections.singletonList; +import static org.awaitility.Awaitility.await; +import static org.dataloader.DataLoaderFactory.newDataLoader; +import static org.dataloader.DataLoaderFactory.newPublisherDataLoader; +import static org.dataloader.DataLoaderFactory.newPublisherDataLoaderWithTry; +import static org.dataloader.DataLoaderOptions.newOptions; +import static org.dataloader.fixtures.TestKit.listFrom; +import static org.dataloader.impl.CompletableFutureKit.cause; +import static org.hamcrest.Matchers.empty; +import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.Matchers.instanceOf; +import static org.hamcrest.Matchers.is; +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertThat; + +public class DataLoaderBatchPublisherTest { + + @Test + public void should_Build_a_really_really_simple_data_loader() { + AtomicBoolean success = new AtomicBoolean(); + DataLoader identityLoader = newPublisherDataLoader(keysAsValues()); + + CompletionStage future1 = identityLoader.load(1); + + future1.thenAccept(value -> { + assertThat(value, equalTo(1)); + success.set(true); + }); + identityLoader.dispatch(); + await().untilAtomic(success, is(true)); + } + + @Test + public void should_Support_loading_multiple_keys_in_one_call() { + AtomicBoolean success = new AtomicBoolean(); + DataLoader identityLoader = newPublisherDataLoader(keysAsValues()); + + CompletionStage> futureAll = identityLoader.loadMany(asList(1, 2)); + futureAll.thenAccept(promisedValues -> { + assertThat(promisedValues.size(), is(2)); + success.set(true); + }); + identityLoader.dispatch(); + await().untilAtomic(success, is(true)); + assertThat(futureAll.toCompletableFuture().join(), equalTo(asList(1, 2))); + } + + @Test + public void should_Resolve_to_empty_list_when_no_keys_supplied() { + AtomicBoolean success = new AtomicBoolean(); + DataLoader identityLoader = newPublisherDataLoader(keysAsValues()); + CompletableFuture> futureEmpty = identityLoader.loadMany(emptyList()); + futureEmpty.thenAccept(promisedValues -> { + assertThat(promisedValues.size(), is(0)); + success.set(true); + }); + identityLoader.dispatch(); + await().untilAtomic(success, is(true)); + assertThat(futureEmpty.join(), empty()); + } + + @Test + public void should_Return_zero_entries_dispatched_when_no_keys_supplied() { + AtomicBoolean success = new AtomicBoolean(); + DataLoader identityLoader = newPublisherDataLoader(keysAsValues()); + CompletableFuture> futureEmpty = identityLoader.loadMany(emptyList()); + futureEmpty.thenAccept(promisedValues -> { + assertThat(promisedValues.size(), is(0)); + success.set(true); + }); + DispatchResult dispatchResult = identityLoader.dispatchWithCounts(); + await().untilAtomic(success, is(true)); + assertThat(dispatchResult.getKeysCount(), equalTo(0)); + } + + @Test + public void should_Batch_multiple_requests() throws ExecutionException, InterruptedException { + List> loadCalls = new ArrayList<>(); + DataLoader identityLoader = idLoader(new DataLoaderOptions(), loadCalls); + + CompletableFuture future1 = identityLoader.load(1); + CompletableFuture future2 = identityLoader.load(2); + identityLoader.dispatch(); + + await().until(() -> future1.isDone() && future2.isDone()); + assertThat(future1.get(), equalTo(1)); + assertThat(future2.get(), equalTo(2)); + assertThat(loadCalls, equalTo(singletonList(asList(1, 2)))); + } + + @Test + public void should_Return_number_of_batched_entries() { + List> loadCalls = new ArrayList<>(); + DataLoader identityLoader = idLoader(new DataLoaderOptions(), loadCalls); + + CompletableFuture future1 = identityLoader.load(1); + CompletableFuture future2 = identityLoader.load(2); + DispatchResult dispatchResult = identityLoader.dispatchWithCounts(); + + await().until(() -> future1.isDone() && future2.isDone()); + assertThat(dispatchResult.getKeysCount(), equalTo(2)); // its two because it's the number dispatched (by key) not the load calls + assertThat(dispatchResult.getPromisedResults().isDone(), equalTo(true)); + } + + @Test + public void should_Coalesce_identical_requests() throws ExecutionException, InterruptedException { + List> loadCalls = new ArrayList<>(); + DataLoader identityLoader = idLoader(new DataLoaderOptions(), loadCalls); + + CompletableFuture future1a = identityLoader.load(1); + CompletableFuture future1b = identityLoader.load(1); + assertThat(future1a, equalTo(future1b)); + identityLoader.dispatch(); + + await().until(future1a::isDone); + assertThat(future1a.get(), equalTo(1)); + assertThat(future1b.get(), equalTo(1)); + assertThat(loadCalls, equalTo(singletonList(singletonList(1)))); + } + + @Test + public void should_Cache_repeated_requests() throws ExecutionException, InterruptedException { + List> loadCalls = new ArrayList<>(); + DataLoader identityLoader = idLoader(new DataLoaderOptions(), loadCalls); + + CompletableFuture future1 = identityLoader.load("A"); + CompletableFuture future2 = identityLoader.load("B"); + identityLoader.dispatch(); + + await().until(() -> future1.isDone() && future2.isDone()); + assertThat(future1.get(), equalTo("A")); + assertThat(future2.get(), equalTo("B")); + assertThat(loadCalls, equalTo(singletonList(asList("A", "B")))); + + CompletableFuture future1a = identityLoader.load("A"); + CompletableFuture future3 = identityLoader.load("C"); + identityLoader.dispatch(); + + await().until(() -> future1a.isDone() && future3.isDone()); + assertThat(future1a.get(), equalTo("A")); + assertThat(future3.get(), equalTo("C")); + assertThat(loadCalls, equalTo(asList(asList("A", "B"), singletonList("C")))); + + CompletableFuture future1b = identityLoader.load("A"); + CompletableFuture future2a = identityLoader.load("B"); + CompletableFuture future3a = identityLoader.load("C"); + identityLoader.dispatch(); + + await().until(() -> future1b.isDone() && future2a.isDone() && future3a.isDone()); + assertThat(future1b.get(), equalTo("A")); + assertThat(future2a.get(), equalTo("B")); + assertThat(future3a.get(), equalTo("C")); + assertThat(loadCalls, equalTo(asList(asList("A", "B"), singletonList("C")))); + } + + @Test + public void should_Not_redispatch_previous_load() throws ExecutionException, InterruptedException { + List> loadCalls = new ArrayList<>(); + DataLoader identityLoader = idLoader(new DataLoaderOptions(), loadCalls); + + CompletableFuture future1 = identityLoader.load("A"); + identityLoader.dispatch(); + + CompletableFuture future2 = identityLoader.load("B"); + identityLoader.dispatch(); + + await().until(() -> future1.isDone() && future2.isDone()); + assertThat(future1.get(), equalTo("A")); + assertThat(future2.get(), equalTo("B")); + assertThat(loadCalls, equalTo(asList(singletonList("A"), singletonList("B")))); + } + + @Test + public void should_Cache_on_redispatch() throws ExecutionException, InterruptedException { + List> loadCalls = new ArrayList<>(); + DataLoader identityLoader = idLoader(new DataLoaderOptions(), loadCalls); + + CompletableFuture future1 = identityLoader.load("A"); + identityLoader.dispatch(); + + CompletableFuture> future2 = identityLoader.loadMany(asList("A", "B")); + identityLoader.dispatch(); + + await().until(() -> future1.isDone() && future2.isDone()); + assertThat(future1.get(), equalTo("A")); + assertThat(future2.get(), equalTo(asList("A", "B"))); + assertThat(loadCalls, equalTo(asList(singletonList("A"), singletonList("B")))); + } + + @Test + public void should_Clear_single_value_in_loader() throws ExecutionException, InterruptedException { + List> loadCalls = new ArrayList<>(); + DataLoader identityLoader = idLoader(new DataLoaderOptions(), loadCalls); + + CompletableFuture future1 = identityLoader.load("A"); + CompletableFuture future2 = identityLoader.load("B"); + identityLoader.dispatch(); + + await().until(() -> future1.isDone() && future2.isDone()); + assertThat(future1.get(), equalTo("A")); + assertThat(future2.get(), equalTo("B")); + assertThat(loadCalls, equalTo(singletonList(asList("A", "B")))); + + // fluency + DataLoader dl = identityLoader.clear("A"); + assertThat(dl, equalTo(identityLoader)); + + CompletableFuture future1a = identityLoader.load("A"); + CompletableFuture future2a = identityLoader.load("B"); + identityLoader.dispatch(); + + await().until(() -> future1a.isDone() && future2a.isDone()); + assertThat(future1a.get(), equalTo("A")); + assertThat(future2a.get(), equalTo("B")); + assertThat(loadCalls, equalTo(asList(asList("A", "B"), singletonList("A")))); + } + + @Test + public void should_Clear_all_values_in_loader() throws ExecutionException, InterruptedException { + List> loadCalls = new ArrayList<>(); + DataLoader identityLoader = idLoader(new DataLoaderOptions(), loadCalls); + + CompletableFuture future1 = identityLoader.load("A"); + CompletableFuture future2 = identityLoader.load("B"); + identityLoader.dispatch(); + + await().until(() -> future1.isDone() && future2.isDone()); + assertThat(future1.get(), equalTo("A")); + assertThat(future2.get(), equalTo("B")); + assertThat(loadCalls, equalTo(singletonList(asList("A", "B")))); + + DataLoader dlFluent = identityLoader.clearAll(); + assertThat(dlFluent, equalTo(identityLoader)); // fluency + + CompletableFuture future1a = identityLoader.load("A"); + CompletableFuture future2a = identityLoader.load("B"); + identityLoader.dispatch(); + + await().until(() -> future1a.isDone() && future2a.isDone()); + assertThat(future1a.get(), equalTo("A")); + assertThat(future2a.get(), equalTo("B")); + assertThat(loadCalls, equalTo(asList(asList("A", "B"), asList("A", "B")))); + } + + @Test + public void should_Allow_priming_the_cache() throws ExecutionException, InterruptedException { + List> loadCalls = new ArrayList<>(); + DataLoader identityLoader = idLoader(new DataLoaderOptions(), loadCalls); + + DataLoader dlFluency = identityLoader.prime("A", "A"); + assertThat(dlFluency, equalTo(identityLoader)); + + CompletableFuture future1 = identityLoader.load("A"); + CompletableFuture future2 = identityLoader.load("B"); + identityLoader.dispatch(); + + await().until(() -> future1.isDone() && future2.isDone()); + assertThat(future1.get(), equalTo("A")); + assertThat(future2.get(), equalTo("B")); + assertThat(loadCalls, equalTo(singletonList(singletonList("B")))); + } + + @Test + public void should_Not_prime_keys_that_already_exist() throws ExecutionException, InterruptedException { + List> loadCalls = new ArrayList<>(); + DataLoader identityLoader = idLoader(new DataLoaderOptions(), loadCalls); + + identityLoader.prime("A", "X"); + + CompletableFuture future1 = identityLoader.load("A"); + CompletableFuture future2 = identityLoader.load("B"); + CompletableFuture> composite = identityLoader.dispatch(); + + await().until(composite::isDone); + assertThat(future1.get(), equalTo("X")); + assertThat(future2.get(), equalTo("B")); + + identityLoader.prime("A", "Y"); + identityLoader.prime("B", "Y"); + + CompletableFuture future1a = identityLoader.load("A"); + CompletableFuture future2a = identityLoader.load("B"); + CompletableFuture> composite2 = identityLoader.dispatch(); + + await().until(composite2::isDone); + assertThat(future1a.get(), equalTo("X")); + assertThat(future2a.get(), equalTo("B")); + assertThat(loadCalls, equalTo(singletonList(singletonList("B")))); + } + + @Test + public void should_Allow_to_forcefully_prime_the_cache() throws ExecutionException, InterruptedException { + List> loadCalls = new ArrayList<>(); + DataLoader identityLoader = idLoader(new DataLoaderOptions(), loadCalls); + + identityLoader.prime("A", "X"); + + CompletableFuture future1 = identityLoader.load("A"); + CompletableFuture future2 = identityLoader.load("B"); + CompletableFuture> composite = identityLoader.dispatch(); + + await().until(composite::isDone); + assertThat(future1.get(), equalTo("X")); + assertThat(future2.get(), equalTo("B")); + + identityLoader.clear("A").prime("A", "Y"); + identityLoader.clear("B").prime("B", "Y"); + + CompletableFuture future1a = identityLoader.load("A"); + CompletableFuture future2a = identityLoader.load("B"); + CompletableFuture> composite2 = identityLoader.dispatch(); + + await().until(composite2::isDone); + assertThat(future1a.get(), equalTo("Y")); + assertThat(future2a.get(), equalTo("Y")); + assertThat(loadCalls, equalTo(singletonList(singletonList("B")))); + } + + @Test + public void should_Allow_priming_the_cache_with_a_future() throws ExecutionException, InterruptedException { + List> loadCalls = new ArrayList<>(); + DataLoader identityLoader = idLoader(new DataLoaderOptions(), loadCalls); + + DataLoader dlFluency = identityLoader.prime("A", CompletableFuture.completedFuture("A")); + assertThat(dlFluency, equalTo(identityLoader)); + + CompletableFuture future1 = identityLoader.load("A"); + CompletableFuture future2 = identityLoader.load("B"); + identityLoader.dispatch(); + + await().until(() -> future1.isDone() && future2.isDone()); + assertThat(future1.get(), equalTo("A")); + assertThat(future2.get(), equalTo("B")); + assertThat(loadCalls, equalTo(singletonList(singletonList("B")))); + } + + @Test + public void should_not_Cache_failed_fetches_on_complete_failure() { + List> loadCalls = new ArrayList<>(); + DataLoader errorLoader = idLoaderBlowsUps(new DataLoaderOptions(), loadCalls); + + CompletableFuture future1 = errorLoader.load(1); + errorLoader.dispatch(); + + await().until(future1::isDone); + assertThat(future1.isCompletedExceptionally(), is(true)); + assertThat(cause(future1), instanceOf(IllegalStateException.class)); + + CompletableFuture future2 = errorLoader.load(1); + errorLoader.dispatch(); + + await().until(future2::isDone); + assertThat(future2.isCompletedExceptionally(), is(true)); + assertThat(cause(future2), instanceOf(IllegalStateException.class)); + assertThat(loadCalls, equalTo(asList(singletonList(1), singletonList(1)))); + } + + @Test + public void should_Resolve_to_error_to_indicate_failure() throws ExecutionException, InterruptedException { + List> loadCalls = new ArrayList<>(); + DataLoader evenLoader = idLoaderOddEvenExceptions(new DataLoaderOptions(), loadCalls); + + CompletableFuture future1 = evenLoader.load(1); + evenLoader.dispatch(); + + await().until(future1::isDone); + assertThat(future1.isCompletedExceptionally(), is(true)); + assertThat(cause(future1), instanceOf(IllegalStateException.class)); + + CompletableFuture future2 = evenLoader.load(2); + evenLoader.dispatch(); + + await().until(future2::isDone); + assertThat(future2.get(), equalTo(2)); + assertThat(loadCalls, equalTo(asList(singletonList(1), singletonList(2)))); + } + + // Accept any kind of key. + + @Test + public void should_Represent_failures_and_successes_simultaneously() throws ExecutionException, InterruptedException { + AtomicBoolean success = new AtomicBoolean(); + List> loadCalls = new ArrayList<>(); + DataLoader evenLoader = idLoaderOddEvenExceptions(new DataLoaderOptions(), loadCalls); + + CompletableFuture future1 = evenLoader.load(1); + CompletableFuture future2 = evenLoader.load(2); + CompletableFuture future3 = evenLoader.load(3); + CompletableFuture future4 = evenLoader.load(4); + CompletableFuture> result = evenLoader.dispatch(); + result.thenAccept(promisedValues -> success.set(true)); + + await().untilAtomic(success, is(true)); + + assertThat(future1.isCompletedExceptionally(), is(true)); + assertThat(cause(future1), instanceOf(IllegalStateException.class)); + assertThat(future2.get(), equalTo(2)); + assertThat(future3.isCompletedExceptionally(), is(true)); + assertThat(future4.get(), equalTo(4)); + + assertThat(loadCalls, equalTo(singletonList(asList(1, 2, 3, 4)))); + } + + // Accepts options + + @Test + public void should_Cache_failed_fetches() { + List> loadCalls = new ArrayList<>(); + DataLoader errorLoader = idLoaderAllExceptions(new DataLoaderOptions(), loadCalls); + + CompletableFuture future1 = errorLoader.load(1); + errorLoader.dispatch(); + + await().until(future1::isDone); + assertThat(future1.isCompletedExceptionally(), is(true)); + assertThat(cause(future1), instanceOf(IllegalStateException.class)); + + CompletableFuture future2 = errorLoader.load(1); + errorLoader.dispatch(); + + await().until(future2::isDone); + assertThat(future2.isCompletedExceptionally(), is(true)); + assertThat(cause(future2), instanceOf(IllegalStateException.class)); + + assertThat(loadCalls, equalTo(singletonList(singletonList(1)))); + } + + @Test + public void should_NOT_Cache_failed_fetches_if_told_not_too() { + DataLoaderOptions options = DataLoaderOptions.newOptions().setCachingExceptionsEnabled(false); + List> loadCalls = new ArrayList<>(); + DataLoader errorLoader = idLoaderAllExceptions(options, loadCalls); + + CompletableFuture future1 = errorLoader.load(1); + errorLoader.dispatch(); + + await().until(future1::isDone); + assertThat(future1.isCompletedExceptionally(), is(true)); + assertThat(cause(future1), instanceOf(IllegalStateException.class)); + + CompletableFuture future2 = errorLoader.load(1); + errorLoader.dispatch(); + + await().until(future2::isDone); + assertThat(future2.isCompletedExceptionally(), is(true)); + assertThat(cause(future2), instanceOf(IllegalStateException.class)); + + assertThat(loadCalls, equalTo(asList(singletonList(1), singletonList(1)))); + } + + + // Accepts object key in custom cacheKey function + + @Test + public void should_Handle_priming_the_cache_with_an_error() { + List> loadCalls = new ArrayList<>(); + DataLoader identityLoader = idLoader(new DataLoaderOptions(), loadCalls); + + identityLoader.prime(1, new IllegalStateException("Error")); + + CompletableFuture future1 = identityLoader.load(1); + identityLoader.dispatch(); + + await().until(future1::isDone); + assertThat(future1.isCompletedExceptionally(), is(true)); + assertThat(cause(future1), instanceOf(IllegalStateException.class)); + assertThat(loadCalls, equalTo(emptyList())); + } + + @Test + public void should_Clear_values_from_cache_after_errors() { + List> loadCalls = new ArrayList<>(); + DataLoader errorLoader = idLoaderBlowsUps(new DataLoaderOptions(), loadCalls); + + CompletableFuture future1 = errorLoader.load(1); + future1.handle((value, t) -> { + if (t != null) { + // Presumably determine if this error is transient, and only clear the cache in that case. + errorLoader.clear(1); + } + return null; + }); + errorLoader.dispatch(); + + await().until(future1::isDone); + assertThat(future1.isCompletedExceptionally(), is(true)); + assertThat(cause(future1), instanceOf(IllegalStateException.class)); + + CompletableFuture future2 = errorLoader.load(1); + future2.handle((value, t) -> { + if (t != null) { + // Again, only do this if you can determine the error is transient. + errorLoader.clear(1); + } + return null; + }); + errorLoader.dispatch(); + + await().until(future2::isDone); + assertThat(future2.isCompletedExceptionally(), is(true)); + assertThat(cause(future2), instanceOf(IllegalStateException.class)); + assertThat(loadCalls, equalTo(asList(singletonList(1), singletonList(1)))); + } + + @Test + public void should_Propagate_error_to_all_loads() { + List> loadCalls = new ArrayList<>(); + DataLoader errorLoader = idLoaderBlowsUps(new DataLoaderOptions(), loadCalls); + + CompletableFuture future1 = errorLoader.load(1); + CompletableFuture future2 = errorLoader.load(2); + errorLoader.dispatch(); + + await().until(future1::isDone); + assertThat(future1.isCompletedExceptionally(), is(true)); + Throwable cause = cause(future1); + assert cause != null; + assertThat(cause, instanceOf(IllegalStateException.class)); + assertThat(cause.getMessage(), equalTo("Error")); + + await().until(future2::isDone); + cause = cause(future2); + assert cause != null; + assertThat(cause.getMessage(), equalTo(cause.getMessage())); + assertThat(loadCalls, equalTo(singletonList(asList(1, 2)))); + } + + @Test + public void should_Accept_objects_as_keys() { + List> loadCalls = new ArrayList<>(); + DataLoader identityLoader = idLoader(new DataLoaderOptions(), loadCalls); + + Object keyA = new Object(); + Object keyB = new Object(); + + // Fetches as expected + + identityLoader.load(keyA); + identityLoader.load(keyB); + + identityLoader.dispatch().thenAccept(promisedValues -> { + assertThat(promisedValues.get(0), equalTo(keyA)); + assertThat(promisedValues.get(1), equalTo(keyB)); + }); + + assertThat(loadCalls.size(), equalTo(1)); + assertThat(loadCalls.get(0).size(), equalTo(2)); + assertThat(loadCalls.get(0).toArray()[0], equalTo(keyA)); + assertThat(loadCalls.get(0).toArray()[1], equalTo(keyB)); + + // Caching + identityLoader.clear(keyA); + //noinspection SuspiciousMethodCalls + loadCalls.remove(keyA); + + identityLoader.load(keyA); + identityLoader.load(keyB); + + identityLoader.dispatch().thenAccept(promisedValues -> { + assertThat(promisedValues.get(0), equalTo(keyA)); + assertThat(identityLoader.getCacheKey(keyB), equalTo(keyB)); + }); + + assertThat(loadCalls.size(), equalTo(2)); + assertThat(loadCalls.get(1).size(), equalTo(1)); + assertThat(loadCalls.get(1).toArray()[0], equalTo(keyA)); + } + + @Test + public void should_Disable_caching() throws ExecutionException, InterruptedException { + List> loadCalls = new ArrayList<>(); + DataLoader identityLoader = + idLoader(newOptions().setCachingEnabled(false), loadCalls); + + CompletableFuture future1 = identityLoader.load("A"); + CompletableFuture future2 = identityLoader.load("B"); + identityLoader.dispatch(); + + await().until(() -> future1.isDone() && future2.isDone()); + assertThat(future1.get(), equalTo("A")); + assertThat(future2.get(), equalTo("B")); + assertThat(loadCalls, equalTo(singletonList(asList("A", "B")))); + + CompletableFuture future1a = identityLoader.load("A"); + CompletableFuture future3 = identityLoader.load("C"); + identityLoader.dispatch(); + + await().until(() -> future1a.isDone() && future3.isDone()); + assertThat(future1a.get(), equalTo("A")); + assertThat(future3.get(), equalTo("C")); + assertThat(loadCalls, equalTo(asList(asList("A", "B"), asList("A", "C")))); + + CompletableFuture future1b = identityLoader.load("A"); + CompletableFuture future2a = identityLoader.load("B"); + CompletableFuture future3a = identityLoader.load("C"); + identityLoader.dispatch(); + + await().until(() -> future1b.isDone() && future2a.isDone() && future3a.isDone()); + assertThat(future1b.get(), equalTo("A")); + assertThat(future2a.get(), equalTo("B")); + assertThat(future3a.get(), equalTo("C")); + assertThat(loadCalls, equalTo(asList(asList("A", "B"), + asList("A", "C"), asList("A", "B", "C")))); + } + + @Test + public void should_work_with_duplicate_keys_when_caching_disabled() throws ExecutionException, InterruptedException { + List> loadCalls = new ArrayList<>(); + DataLoader identityLoader = + idLoader(newOptions().setCachingEnabled(false), loadCalls); + + CompletableFuture future1 = identityLoader.load("A"); + CompletableFuture future2 = identityLoader.load("B"); + CompletableFuture future3 = identityLoader.load("A"); + identityLoader.dispatch(); + + await().until(() -> future1.isDone() && future2.isDone() && future3.isDone()); + assertThat(future1.get(), equalTo("A")); + assertThat(future2.get(), equalTo("B")); + assertThat(future3.get(), equalTo("A")); + assertThat(loadCalls, equalTo(singletonList(asList("A", "B", "A")))); + } + + @Test + public void should_work_with_duplicate_keys_when_caching_enabled() throws ExecutionException, InterruptedException { + List> loadCalls = new ArrayList<>(); + DataLoader identityLoader = + idLoader(newOptions().setCachingEnabled(true), loadCalls); + + CompletableFuture future1 = identityLoader.load("A"); + CompletableFuture future2 = identityLoader.load("B"); + CompletableFuture future3 = identityLoader.load("A"); + identityLoader.dispatch(); + + await().until(() -> future1.isDone() && future2.isDone() && future3.isDone()); + assertThat(future1.get(), equalTo("A")); + assertThat(future2.get(), equalTo("B")); + assertThat(future3.get(), equalTo("A")); + assertThat(loadCalls, equalTo(singletonList(asList("A", "B")))); + } + + // It is resilient to job queue ordering + + @Test + public void should_Accept_objects_with_a_complex_key() throws ExecutionException, InterruptedException { + List> loadCalls = new ArrayList<>(); + DataLoaderOptions options = newOptions().setCacheKeyFunction(getJsonObjectCacheMapFn()); + DataLoader identityLoader = idLoader(options, loadCalls); + + JsonObject key1 = new JsonObject().put("id", 123); + JsonObject key2 = new JsonObject().put("id", 123); + + CompletableFuture future1 = identityLoader.load(key1); + CompletableFuture future2 = identityLoader.load(key2); + identityLoader.dispatch(); + + await().until(() -> future1.isDone() && future2.isDone()); + assertThat(loadCalls, equalTo(singletonList(singletonList(key1)))); + assertThat(future1.get(), equalTo(key1)); + assertThat(future2.get(), equalTo(key1)); + } + + // Helper methods + + @Test + public void should_Clear_objects_with_complex_key() throws ExecutionException, InterruptedException { + List> loadCalls = new ArrayList<>(); + DataLoaderOptions options = newOptions().setCacheKeyFunction(getJsonObjectCacheMapFn()); + DataLoader identityLoader = idLoader(options, loadCalls); + + JsonObject key1 = new JsonObject().put("id", 123); + JsonObject key2 = new JsonObject().put("id", 123); + + CompletableFuture future1 = identityLoader.load(key1); + identityLoader.dispatch(); + + await().until(future1::isDone); + identityLoader.clear(key2); // clear equivalent object key + + CompletableFuture future2 = identityLoader.load(key1); + identityLoader.dispatch(); + + await().until(future2::isDone); + assertThat(loadCalls, equalTo(asList(singletonList(key1), singletonList(key1)))); + assertThat(future1.get(), equalTo(key1)); + assertThat(future2.get(), equalTo(key1)); + } + + @Test + public void should_Accept_objects_with_different_order_of_keys() throws ExecutionException, InterruptedException { + List> loadCalls = new ArrayList<>(); + DataLoaderOptions options = newOptions().setCacheKeyFunction(getJsonObjectCacheMapFn()); + DataLoader identityLoader = idLoader(options, loadCalls); + + JsonObject key1 = new JsonObject().put("a", 123).put("b", 321); + JsonObject key2 = new JsonObject().put("b", 321).put("a", 123); + + // Fetches as expected + + CompletableFuture future1 = identityLoader.load(key1); + CompletableFuture future2 = identityLoader.load(key2); + identityLoader.dispatch(); + + await().until(() -> future1.isDone() && future2.isDone()); + assertThat(loadCalls, equalTo(singletonList(singletonList(key1)))); + assertThat(loadCalls.size(), equalTo(1)); + assertThat(future1.get(), equalTo(key1)); + assertThat(future2.get(), equalTo(key2)); + } + + @Test + public void should_Allow_priming_the_cache_with_an_object_key() throws ExecutionException, InterruptedException { + List> loadCalls = new ArrayList<>(); + DataLoaderOptions options = newOptions().setCacheKeyFunction(getJsonObjectCacheMapFn()); + DataLoader identityLoader = idLoader(options, loadCalls); + + JsonObject key1 = new JsonObject().put("id", 123); + JsonObject key2 = new JsonObject().put("id", 123); + + identityLoader.prime(key1, key1); + + CompletableFuture future1 = identityLoader.load(key1); + CompletableFuture future2 = identityLoader.load(key2); + identityLoader.dispatch(); + + await().until(() -> future1.isDone() && future2.isDone()); + assertThat(loadCalls, equalTo(emptyList())); + assertThat(future1.get(), equalTo(key1)); + assertThat(future2.get(), equalTo(key1)); + } + + @Test + public void should_Accept_a_custom_cache_map_implementation() throws ExecutionException, InterruptedException { + CustomCacheMap customMap = new CustomCacheMap(); + List> loadCalls = new ArrayList<>(); + DataLoaderOptions options = newOptions().setCacheMap(customMap); + DataLoader identityLoader = idLoader(options, loadCalls); + + // Fetches as expected + + CompletableFuture future1 = identityLoader.load("a"); + CompletableFuture future2 = identityLoader.load("b"); + CompletableFuture> composite = identityLoader.dispatch(); + + await().until(composite::isDone); + assertThat(future1.get(), equalTo("a")); + assertThat(future2.get(), equalTo("b")); + + assertThat(loadCalls, equalTo(singletonList(asList("a", "b")))); + assertArrayEquals(customMap.stash.keySet().toArray(), asList("a", "b").toArray()); + + CompletableFuture future3 = identityLoader.load("c"); + CompletableFuture future2a = identityLoader.load("b"); + composite = identityLoader.dispatch(); + + await().until(composite::isDone); + assertThat(future3.get(), equalTo("c")); + assertThat(future2a.get(), equalTo("b")); + + assertThat(loadCalls, equalTo(asList(asList("a", "b"), singletonList("c")))); + assertArrayEquals(customMap.stash.keySet().toArray(), asList("a", "b", "c").toArray()); + + // Supports clear + + identityLoader.clear("b"); + assertArrayEquals(customMap.stash.keySet().toArray(), asList("a", "c").toArray()); + + CompletableFuture future2b = identityLoader.load("b"); + composite = identityLoader.dispatch(); + + await().until(composite::isDone); + assertThat(future2b.get(), equalTo("b")); + assertThat(loadCalls, equalTo(asList(asList("a", "b"), + singletonList("c"), singletonList("b")))); + assertArrayEquals(customMap.stash.keySet().toArray(), asList("a", "c", "b").toArray()); + + // Supports clear all + + identityLoader.clearAll(); + assertArrayEquals(customMap.stash.keySet().toArray(), emptyList().toArray()); + } + + @Test + public void should_degrade_gracefully_if_cache_get_throws() { + CacheMap cache = new ThrowingCacheMap(); + DataLoaderOptions options = newOptions().setCachingEnabled(true).setCacheMap(cache); + List> loadCalls = new ArrayList<>(); + DataLoader identityLoader = idLoader(options, loadCalls); + + assertThat(identityLoader.getIfPresent("a"), equalTo(Optional.empty())); + + CompletableFuture future = identityLoader.load("a"); + identityLoader.dispatch(); + assertThat(future.join(), equalTo("a")); + } + + @Test + public void batching_disabled_should_dispatch_immediately() { + List> loadCalls = new ArrayList<>(); + DataLoaderOptions options = newOptions().setBatchingEnabled(false); + DataLoader identityLoader = idLoader(options, loadCalls); + + CompletableFuture fa = identityLoader.load("A"); + CompletableFuture fb = identityLoader.load("B"); + + // caching is on still + CompletableFuture fa1 = identityLoader.load("A"); + CompletableFuture fb1 = identityLoader.load("B"); + + List values = CompletableFutureKit.allOf(asList(fa, fb, fa1, fb1)).join(); + + assertThat(fa.join(), equalTo("A")); + assertThat(fb.join(), equalTo("B")); + assertThat(fa1.join(), equalTo("A")); + assertThat(fb1.join(), equalTo("B")); + + assertThat(values, equalTo(asList("A", "B", "A", "B"))); + + assertThat(loadCalls, equalTo(asList( + singletonList("A"), + singletonList("B")))); + + } + + @Test + public void batching_disabled_and_caching_disabled_should_dispatch_immediately_and_forget() { + List> loadCalls = new ArrayList<>(); + DataLoaderOptions options = newOptions().setBatchingEnabled(false).setCachingEnabled(false); + DataLoader identityLoader = idLoader(options, loadCalls); + + CompletableFuture fa = identityLoader.load("A"); + CompletableFuture fb = identityLoader.load("B"); + + // caching is off + CompletableFuture fa1 = identityLoader.load("A"); + CompletableFuture fb1 = identityLoader.load("B"); + + List values = CompletableFutureKit.allOf(asList(fa, fb, fa1, fb1)).join(); + + assertThat(fa.join(), equalTo("A")); + assertThat(fb.join(), equalTo("B")); + assertThat(fa1.join(), equalTo("A")); + assertThat(fb1.join(), equalTo("B")); + + assertThat(values, equalTo(asList("A", "B", "A", "B"))); + + assertThat(loadCalls, equalTo(asList( + singletonList("A"), + singletonList("B"), + singletonList("A"), + singletonList("B") + ))); + + } + + @Test + public void batches_multiple_requests_with_max_batch_size() { + List> loadCalls = new ArrayList<>(); + DataLoader identityLoader = idLoader(newOptions().setMaxBatchSize(2), loadCalls); + + CompletableFuture f1 = identityLoader.load(1); + CompletableFuture f2 = identityLoader.load(2); + CompletableFuture f3 = identityLoader.load(3); + + identityLoader.dispatch(); + + CompletableFuture.allOf(f1, f2, f3).join(); + + assertThat(f1.join(), equalTo(1)); + assertThat(f2.join(), equalTo(2)); + assertThat(f3.join(), equalTo(3)); + + assertThat(loadCalls, equalTo(asList(asList(1, 2), singletonList(3)))); + + } + + @Test + public void can_split_max_batch_sizes_correctly() { + List> loadCalls = new ArrayList<>(); + DataLoader identityLoader = idLoader(newOptions().setMaxBatchSize(5), loadCalls); + + for (int i = 0; i < 21; i++) { + identityLoader.load(i); + } + List> expectedCalls = new ArrayList<>(); + expectedCalls.add(listFrom(0, 5)); + expectedCalls.add(listFrom(5, 10)); + expectedCalls.add(listFrom(10, 15)); + expectedCalls.add(listFrom(15, 20)); + expectedCalls.add(listFrom(20, 21)); + + List result = identityLoader.dispatch().join(); + + assertThat(result, equalTo(listFrom(0, 21))); + assertThat(loadCalls, equalTo(expectedCalls)); + + } + + @Test + public void should_Batch_loads_occurring_within_futures() { + List> loadCalls = new ArrayList<>(); + DataLoader identityLoader = idLoader(newOptions(), loadCalls); + + Supplier nullValue = () -> null; + + AtomicBoolean v4Called = new AtomicBoolean(); + + CompletableFuture.supplyAsync(nullValue).thenAccept(v1 -> { + identityLoader.load("a"); + CompletableFuture.supplyAsync(nullValue).thenAccept(v2 -> { + identityLoader.load("b"); + CompletableFuture.supplyAsync(nullValue).thenAccept(v3 -> { + identityLoader.load("c"); + CompletableFuture.supplyAsync(nullValue).thenAccept( + v4 -> { + identityLoader.load("d"); + v4Called.set(true); + }); + }); + }); + }); + + await().untilTrue(v4Called); + + identityLoader.dispatchAndJoin(); + + assertThat(loadCalls, equalTo( + singletonList(asList("a", "b", "c", "d")))); + } + + @Test + public void can_call_a_loader_from_a_loader() throws Exception { + List> deepLoadCalls = new ArrayList<>(); + DataLoader deepLoader = newDataLoader(keys -> { + deepLoadCalls.add(keys); + return CompletableFuture.completedFuture(keys); + }); + + List> aLoadCalls = new ArrayList<>(); + DataLoader aLoader = newDataLoader(keys -> { + aLoadCalls.add(keys); + return deepLoader.loadMany(keys); + }); + + List> bLoadCalls = new ArrayList<>(); + DataLoader bLoader = newDataLoader(keys -> { + bLoadCalls.add(keys); + return deepLoader.loadMany(keys); + }); + + CompletableFuture a1 = aLoader.load("A1"); + CompletableFuture a2 = aLoader.load("A2"); + CompletableFuture b1 = bLoader.load("B1"); + CompletableFuture b2 = bLoader.load("B2"); + + CompletableFuture.allOf( + aLoader.dispatch(), + deepLoader.dispatch(), + bLoader.dispatch(), + deepLoader.dispatch() + ).join(); + + assertThat(a1.get(), equalTo("A1")); + assertThat(a2.get(), equalTo("A2")); + assertThat(b1.get(), equalTo("B1")); + assertThat(b2.get(), equalTo("B2")); + + assertThat(aLoadCalls, equalTo( + singletonList(asList("A1", "A2")))); + + assertThat(bLoadCalls, equalTo( + singletonList(asList("B1", "B2")))); + + assertThat(deepLoadCalls, equalTo( + asList(asList("A1", "A2"), asList("B1", "B2")))); + } + + @Test + public void should_allow_composition_of_data_loader_calls() { + UserManager userManager = new UserManager(); + + BatchLoader userBatchLoader = userIds -> CompletableFuture + .supplyAsync(() -> userIds + .stream() + .map(userManager::loadUserById) + .collect(Collectors.toList())); + DataLoader userLoader = newDataLoader(userBatchLoader); + + AtomicBoolean gandalfCalled = new AtomicBoolean(false); + AtomicBoolean sarumanCalled = new AtomicBoolean(false); + + userLoader.load(1L) + .thenAccept(user -> userLoader.load(user.getInvitedByID()) + .thenAccept(invitedBy -> { + gandalfCalled.set(true); + assertThat(invitedBy.getName(), equalTo("Manwë")); + })); + + userLoader.load(2L) + .thenAccept(user -> userLoader.load(user.getInvitedByID()) + .thenAccept(invitedBy -> { + sarumanCalled.set(true); + assertThat(invitedBy.getName(), equalTo("Aulë")); + })); + + List allResults = userLoader.dispatchAndJoin(); + + await().untilTrue(gandalfCalled); + await().untilTrue(sarumanCalled); + + assertThat(allResults.size(), equalTo(4)); + } + + private static CacheKey getJsonObjectCacheMapFn() { + return key -> key.stream() + .map(entry -> entry.getKey() + ":" + entry.getValue()) + .sorted() + .collect(Collectors.joining()); + } + + private static DataLoader idLoader(DataLoaderOptions options, List> loadCalls) { + return newPublisherDataLoader((BatchPublisher) (keys, subscriber) -> { + loadCalls.add(new ArrayList<>(keys)); + Flux.fromIterable(keys).subscribe(subscriber); + }, options); + } + + private static DataLoader idLoaderBlowsUps( + DataLoaderOptions options, List> loadCalls) { + return newPublisherDataLoader((BatchPublisher) (keys, subscriber) -> { + loadCalls.add(new ArrayList<>(keys)); + Flux.error(new IllegalStateException("Error")).subscribe(subscriber); + }, options); + } + + private static DataLoader idLoaderAllExceptions( + DataLoaderOptions options, List> loadCalls) { + return newPublisherDataLoaderWithTry((BatchPublisher>) (keys, subscriber) -> { + loadCalls.add(new ArrayList<>(keys)); + Stream> failures = keys.stream().map(k -> Try.failed(new IllegalStateException("Error"))); + Flux.fromStream(failures).subscribe(subscriber); + }, options); + } + + private static DataLoader idLoaderOddEvenExceptions( + DataLoaderOptions options, List> loadCalls) { + return newPublisherDataLoaderWithTry((BatchPublisher>) (keys, subscriber) -> { + loadCalls.add(new ArrayList<>(keys)); + + List> errors = new ArrayList<>(); + for (Integer key : keys) { + if (key % 2 == 0) { + errors.add(Try.succeeded(key)); + } else { + errors.add(Try.failed(new IllegalStateException("Error"))); + } + } + Flux.fromIterable(errors).subscribe(subscriber); + }, options); + } + + private static BatchPublisher keysAsValues() { + return (keys, subscriber) -> Flux.fromIterable(keys).subscribe(subscriber); + } + + private static class ThrowingCacheMap extends CustomCacheMap { + @Override + public CompletableFuture get(String key) { + throw new RuntimeException("Cache implementation failed."); + } + } +} diff --git a/src/test/java/org/dataloader/DataLoaderMappedBatchPublisherTest.java b/src/test/java/org/dataloader/DataLoaderMappedBatchPublisherTest.java new file mode 100644 index 0000000..c16c58f --- /dev/null +++ b/src/test/java/org/dataloader/DataLoaderMappedBatchPublisherTest.java @@ -0,0 +1,175 @@ +package org.dataloader; + +import org.junit.Test; +import reactor.core.publisher.Flux; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.atomic.AtomicInteger; + +import static java.util.Arrays.asList; +import static java.util.Collections.singletonList; +import static org.awaitility.Awaitility.await; +import static org.dataloader.DataLoaderFactory.newMappedPublisherDataLoader; +import static org.dataloader.DataLoaderOptions.newOptions; +import static org.dataloader.fixtures.TestKit.listFrom; +import static org.dataloader.impl.CompletableFutureKit.cause; +import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.Matchers.instanceOf; +import static org.hamcrest.Matchers.is; +import static org.junit.Assert.assertThat; + +public class DataLoaderMappedBatchPublisherTest { + + MappedBatchPublisher evensOnlyMappedBatchLoader = (keys, subscriber) -> { + Map mapOfResults = new HashMap<>(); + + AtomicInteger index = new AtomicInteger(); + keys.forEach(k -> { + int i = index.getAndIncrement(); + if (i % 2 == 0) { + mapOfResults.put(k, k); + } + }); + Flux.fromIterable(mapOfResults.entrySet()).subscribe(subscriber); + }; + + private static DataLoader idMapLoader(DataLoaderOptions options, List> loadCalls) { + MappedBatchPublisher kvBatchLoader = (keys, subscriber) -> { + loadCalls.add(new ArrayList<>(keys)); + Map map = new HashMap<>(); + //noinspection unchecked + keys.forEach(k -> map.put(k, (V) k)); + Flux.fromIterable(map.entrySet()).subscribe(subscriber); + }; + return DataLoaderFactory.newMappedPublisherDataLoader(kvBatchLoader, options); + } + + private static DataLoader idMapLoaderBlowsUps( + DataLoaderOptions options, List> loadCalls) { + return newMappedPublisherDataLoader((MappedBatchPublisher) (keys, subscriber) -> { + loadCalls.add(new ArrayList<>(keys)); + Flux.>error(new IllegalStateException("Error")).subscribe(subscriber); + }, options); + } + + + @Test + public void basic_map_batch_loading() { + DataLoader loader = DataLoaderFactory.newMappedPublisherDataLoader(evensOnlyMappedBatchLoader); + + loader.load("A"); + loader.load("B"); + loader.loadMany(asList("C", "D")); + + List results = loader.dispatchAndJoin(); + + assertThat(results.size(), equalTo(4)); + assertThat(results, equalTo(asList("A", null, "C", null))); + } + + @Test + public void should_map_Batch_multiple_requests() throws ExecutionException, InterruptedException { + List> loadCalls = new ArrayList<>(); + DataLoader identityLoader = idMapLoader(new DataLoaderOptions(), loadCalls); + + CompletableFuture future1 = identityLoader.load(1); + CompletableFuture future2 = identityLoader.load(2); + identityLoader.dispatch(); + + await().until(() -> future1.isDone() && future2.isDone()); + assertThat(future1.get(), equalTo(1)); + assertThat(future2.get(), equalTo(2)); + assertThat(loadCalls, equalTo(singletonList(asList(1, 2)))); + } + + @Test + public void can_split_max_batch_sizes_correctly() { + List> loadCalls = new ArrayList<>(); + DataLoader identityLoader = idMapLoader(newOptions().setMaxBatchSize(5), loadCalls); + + for (int i = 0; i < 21; i++) { + identityLoader.load(i); + } + List> expectedCalls = new ArrayList<>(); + expectedCalls.add(listFrom(0, 5)); + expectedCalls.add(listFrom(5, 10)); + expectedCalls.add(listFrom(10, 15)); + expectedCalls.add(listFrom(15, 20)); + expectedCalls.add(listFrom(20, 21)); + + List result = identityLoader.dispatch().join(); + + assertThat(result, equalTo(listFrom(0, 21))); + assertThat(loadCalls, equalTo(expectedCalls)); + } + + @Test + public void should_Propagate_error_to_all_loads() { + List> loadCalls = new ArrayList<>(); + DataLoader errorLoader = idMapLoaderBlowsUps(new DataLoaderOptions(), loadCalls); + + CompletableFuture future1 = errorLoader.load(1); + CompletableFuture future2 = errorLoader.load(2); + errorLoader.dispatch(); + + await().until(future1::isDone); + + assertThat(future1.isCompletedExceptionally(), is(true)); + Throwable cause = cause(future1); + assert cause != null; + assertThat(cause, instanceOf(IllegalStateException.class)); + assertThat(cause.getMessage(), equalTo("Error")); + + await().until(future2::isDone); + cause = cause(future2); + assert cause != null; + assertThat(cause.getMessage(), equalTo(cause.getMessage())); + + assertThat(loadCalls, equalTo(singletonList(asList(1, 2)))); + } + + @Test + public void should_work_with_duplicate_keys_when_caching_disabled() throws ExecutionException, InterruptedException { + List> loadCalls = new ArrayList<>(); + DataLoader identityLoader = + idMapLoader(newOptions().setCachingEnabled(false), loadCalls); + + CompletableFuture future1 = identityLoader.load("A"); + CompletableFuture future2 = identityLoader.load("B"); + CompletableFuture future3 = identityLoader.load("A"); + identityLoader.dispatch(); + + await().until(() -> future1.isDone() && future2.isDone() && future3.isDone()); + assertThat(future1.get(), equalTo("A")); + assertThat(future2.get(), equalTo("B")); + assertThat(future3.get(), equalTo("A")); + + // the map batch functions use a set of keys as input and hence remove duplicates unlike list variant + assertThat(loadCalls, equalTo(singletonList(asList("A", "B")))); + } + + @Test + public void should_work_with_duplicate_keys_when_caching_enabled() throws ExecutionException, InterruptedException { + List> loadCalls = new ArrayList<>(); + DataLoader identityLoader = + idMapLoader(newOptions().setCachingEnabled(true), loadCalls); + + CompletableFuture future1 = identityLoader.load("A"); + CompletableFuture future2 = identityLoader.load("B"); + CompletableFuture future3 = identityLoader.load("A"); + identityLoader.dispatch(); + + await().until(() -> future1.isDone() && future2.isDone() && future3.isDone()); + assertThat(future1.get(), equalTo("A")); + assertThat(future2.get(), equalTo("B")); + assertThat(future3.get(), equalTo("A")); + assertThat(loadCalls, equalTo(singletonList(asList("A", "B")))); + } + +} diff --git a/src/test/java/org/dataloader/scheduler/BatchLoaderSchedulerTest.java b/src/test/java/org/dataloader/scheduler/BatchLoaderSchedulerTest.java index beb7c18..e9c43f8 100644 --- a/src/test/java/org/dataloader/scheduler/BatchLoaderSchedulerTest.java +++ b/src/test/java/org/dataloader/scheduler/BatchLoaderSchedulerTest.java @@ -36,6 +36,11 @@ public CompletionStage> scheduleBatchLoader(ScheduledBatchLoaderC public CompletionStage> scheduleMappedBatchLoader(ScheduledMappedBatchLoaderCall scheduledCall, List keys, BatchLoaderEnvironment environment) { return scheduledCall.invoke(); } + + @Override + public void scheduleBatchPublisher(ScheduledBatchPublisherCall scheduledCall, List keys, BatchLoaderEnvironment environment) { + scheduledCall.invoke(); + } }; private BatchLoaderScheduler delayedScheduling(int ms) { @@ -56,6 +61,12 @@ public CompletionStage> scheduleMappedBatchLoader(ScheduledMapp return scheduledCall.invoke(); }).thenCompose(Function.identity()); } + + @Override + public void scheduleBatchPublisher(ScheduledBatchPublisherCall scheduledCall, List keys, BatchLoaderEnvironment environment) { + snooze(ms); + scheduledCall.invoke(); + } }; } @@ -139,6 +150,15 @@ public CompletionStage> scheduleMappedBatchLoader(ScheduledMapp return scheduledCall.invoke(); }).thenCompose(Function.identity()); } + + @Override + public void scheduleBatchPublisher(ScheduledBatchPublisherCall scheduledCall, List keys, BatchLoaderEnvironment environment) { + CompletableFuture.supplyAsync(() -> { + snooze(10); + scheduledCall.invoke(); + return null; + }); + } }; DataLoaderOptions options = DataLoaderOptions.newOptions().setBatchLoaderScheduler(funkyScheduler);