From 22413bed249295455c53830b438fa96ff6c19551 Mon Sep 17 00:00:00 2001 From: Harrison Cole Date: Thu, 12 Dec 2024 15:30:57 +1100 Subject: [PATCH 1/2] `DataLoader` and `CompletableFutureKit`: add support for `loadMany` given a Map of keys and context objects. --- src/main/java/org/dataloader/DataLoader.java | 34 ++++++++++++++++--- .../dataloader/impl/CompletableFutureKit.java | 15 +++++++- 2 files changed, 44 insertions(+), 5 deletions(-) diff --git a/src/main/java/org/dataloader/DataLoader.java b/src/main/java/org/dataloader/DataLoader.java index 1e4ce7d..dc4b726 100644 --- a/src/main/java/org/dataloader/DataLoader.java +++ b/src/main/java/org/dataloader/DataLoader.java @@ -25,10 +25,7 @@ import java.time.Clock; import java.time.Duration; import java.time.Instant; -import java.util.ArrayList; -import java.util.Collections; -import java.util.List; -import java.util.Optional; +import java.util.*; import java.util.concurrent.CompletableFuture; import java.util.function.BiConsumer; @@ -574,6 +571,35 @@ public CompletableFuture> loadMany(List keys, List keyContext } } + /** + * Requests to load the map of data provided by the specified keys asynchronously, and returns a composite future + * of the resulting values. + *

+ * If batching is enabled (the default), you'll have to call {@link DataLoader#dispatch()} at a later stage to + * start batch execution. If you forget this call the future will never be completed (unless already completed, + * and returned from cache). + *

+ * The key context object may be useful in the batch loader interfaces such as {@link org.dataloader.BatchLoaderWithContext} or + * {@link org.dataloader.MappedBatchLoaderWithContext} to help retrieve data. + * + * @param keysAndContexts the map of keys to their respective contexts + * + * @return the composite future of the map of keys and values + */ + public CompletableFuture> loadMany(Map keysAndContexts) { + nonNull(keysAndContexts); + + synchronized (this) { + Map> collect = new HashMap<>(keysAndContexts.size()); + for (Map.Entry entry : keysAndContexts.entrySet()) { + K key = entry.getKey(); + Object keyContext = entry.getValue(); + collect.put(key, load(key, keyContext)); + } + return CompletableFutureKit.allOf(collect); + } + } + /** * Dispatches the queued load requests to the batch execution function and returns a promise of the result. *

diff --git a/src/main/java/org/dataloader/impl/CompletableFutureKit.java b/src/main/java/org/dataloader/impl/CompletableFutureKit.java index 2b94d10..ebc35ec 100644 --- a/src/main/java/org/dataloader/impl/CompletableFutureKit.java +++ b/src/main/java/org/dataloader/impl/CompletableFutureKit.java @@ -3,8 +3,10 @@ import org.dataloader.annotations.Internal; import java.util.List; +import java.util.Map; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutionException; +import java.util.stream.Collectors; import static java.util.stream.Collectors.toList; @@ -48,10 +50,21 @@ public static boolean failed(CompletableFuture future) { } public static CompletableFuture> allOf(List> cfs) { - return CompletableFuture.allOf(cfs.toArray(new CompletableFuture[0])) + return CompletableFuture.allOf(cfs.toArray(CompletableFuture[]::new)) .thenApply(v -> cfs.stream() .map(CompletableFuture::join) .collect(toList()) ); } + + public static CompletableFuture> allOf(Map> cfs) { + return CompletableFuture.allOf(cfs.values().toArray(CompletableFuture[]::new)) + .thenApply(v -> cfs.entrySet().stream() + .collect( + Collectors.toMap( + Map.Entry::getKey, + task -> task.getValue().join()) + ) + ); + } } From 2f437a75dac6a68de201e66d5f4f294e8a6b8589 Mon Sep 17 00:00:00 2001 From: Harrison Cole Date: Thu, 12 Dec 2024 17:16:55 +1100 Subject: [PATCH 2/2] `DataLoader` and `CompletableFutureKit`: add test cases for `loadMany` given a Map of keys and context objects. --- .../DataLoaderBatchLoaderEnvironmentTest.java | 56 ++++++++--- .../org/dataloader/DataLoaderStatsTest.java | 12 ++- .../java/org/dataloader/DataLoaderTest.java | 93 +++++++++++++++---- 3 files changed, 124 insertions(+), 37 deletions(-) diff --git a/src/test/java/org/dataloader/DataLoaderBatchLoaderEnvironmentTest.java b/src/test/java/org/dataloader/DataLoaderBatchLoaderEnvironmentTest.java index 6d3010e..90adbc5 100644 --- a/src/test/java/org/dataloader/DataLoaderBatchLoaderEnvironmentTest.java +++ b/src/test/java/org/dataloader/DataLoaderBatchLoaderEnvironmentTest.java @@ -2,10 +2,7 @@ import org.junit.jupiter.api.Test; -import java.util.ArrayList; -import java.util.HashMap; -import java.util.List; -import java.util.Map; +import java.util.*; import java.util.concurrent.CompletableFuture; import java.util.concurrent.atomic.AtomicInteger; import java.util.stream.Collectors; @@ -50,10 +47,14 @@ public void context_is_passed_to_batch_loader_function() { loader.load("A"); loader.load("B"); loader.loadMany(asList("C", "D")); + Map keysAndContexts = new LinkedHashMap<>(); + keysAndContexts.put("E", null); + keysAndContexts.put("F", null); + loader.loadMany(keysAndContexts); List results = loader.dispatchAndJoin(); - assertThat(results, equalTo(asList("A-ctx", "B-ctx", "C-ctx", "D-ctx"))); + assertThat(results, equalTo(asList("A-ctx", "B-ctx", "C-ctx", "D-ctx", "E-ctx", "F-ctx"))); } @Test @@ -66,10 +67,14 @@ public void key_contexts_are_passed_to_batch_loader_function() { loader.load("A", "aCtx"); loader.load("B", "bCtx"); loader.loadMany(asList("C", "D"), asList("cCtx", "dCtx")); + Map keysAndContexts = new LinkedHashMap<>(); + keysAndContexts.put("E", "eCtx"); + keysAndContexts.put("F", "fCtx"); + loader.loadMany(keysAndContexts); List results = loader.dispatchAndJoin(); - assertThat(results, equalTo(asList("A-ctx-m:aCtx-l:aCtx", "B-ctx-m:bCtx-l:bCtx", "C-ctx-m:cCtx-l:cCtx", "D-ctx-m:dCtx-l:dCtx"))); + assertThat(results, equalTo(asList("A-ctx-m:aCtx-l:aCtx", "B-ctx-m:bCtx-l:bCtx", "C-ctx-m:cCtx-l:cCtx", "D-ctx-m:dCtx-l:dCtx", "E-ctx-m:eCtx-l:eCtx", "F-ctx-m:fCtx-l:fCtx"))); } @Test @@ -82,12 +87,17 @@ public void key_contexts_are_passed_to_batch_loader_function_when_batching_disab CompletableFuture aLoad = loader.load("A", "aCtx"); CompletableFuture bLoad = loader.load("B", "bCtx"); - CompletableFuture> canDLoad = loader.loadMany(asList("C", "D"), asList("cCtx", "dCtx")); + CompletableFuture> cAndDLoad = loader.loadMany(asList("C", "D"), asList("cCtx", "dCtx")); + Map keysAndContexts = new LinkedHashMap<>(); + keysAndContexts.put("E", "eCtx"); + keysAndContexts.put("F", "fCtx"); + CompletableFuture> eAndFLoad = loader.loadMany(keysAndContexts); List results = new ArrayList<>(asList(aLoad.join(), bLoad.join())); - results.addAll(canDLoad.join()); + results.addAll(cAndDLoad.join()); + results.addAll(eAndFLoad.join().values()); - assertThat(results, equalTo(asList("A-ctx-m:aCtx-l:aCtx", "B-ctx-m:bCtx-l:bCtx", "C-ctx-m:cCtx-l:cCtx", "D-ctx-m:dCtx-l:dCtx"))); + assertThat(results, equalTo(asList("A-ctx-m:aCtx-l:aCtx", "B-ctx-m:bCtx-l:bCtx", "C-ctx-m:cCtx-l:cCtx", "D-ctx-m:dCtx-l:dCtx", "E-ctx-m:eCtx-l:eCtx", "F-ctx-m:fCtx-l:fCtx"))); } @Test @@ -101,9 +111,14 @@ public void missing_key_contexts_are_passed_to_batch_loader_function() { loader.load("B"); loader.loadMany(asList("C", "D"), singletonList("cCtx")); + Map keysAndContexts = new LinkedHashMap<>(); + keysAndContexts.put("E", "eCtx"); + keysAndContexts.put("F", null); + loader.loadMany(keysAndContexts); + List results = loader.dispatchAndJoin(); - assertThat(results, equalTo(asList("A-ctx-m:aCtx-l:aCtx", "B-ctx-m:null-l:null", "C-ctx-m:cCtx-l:cCtx", "D-ctx-m:null-l:null"))); + assertThat(results, equalTo(asList("A-ctx-m:aCtx-l:aCtx", "B-ctx-m:null-l:null", "C-ctx-m:cCtx-l:cCtx", "D-ctx-m:null-l:null", "E-ctx-m:eCtx-l:eCtx", "F-ctx-m:null-l:null"))); } @Test @@ -125,9 +140,14 @@ public void context_is_passed_to_map_batch_loader_function() { loader.load("B"); loader.loadMany(asList("C", "D"), singletonList("cCtx")); + Map keysAndContexts = new LinkedHashMap<>(); + keysAndContexts.put("E", "eCtx"); + keysAndContexts.put("F", null); + loader.loadMany(keysAndContexts); + List results = loader.dispatchAndJoin(); - assertThat(results, equalTo(asList("A-ctx-aCtx", "B-ctx-null", "C-ctx-cCtx", "D-ctx-null"))); + assertThat(results, equalTo(asList("A-ctx-aCtx", "B-ctx-null", "C-ctx-cCtx", "D-ctx-null", "E-ctx-eCtx", "F-ctx-null"))); } @Test @@ -142,9 +162,14 @@ public void null_is_passed_as_context_if_you_do_nothing() { loader.load("B"); loader.loadMany(asList("C", "D")); + Map keysAndContexts = new LinkedHashMap<>(); + keysAndContexts.put("E", null); + keysAndContexts.put("F", null); + loader.loadMany(keysAndContexts); + List results = loader.dispatchAndJoin(); - assertThat(results, equalTo(asList("A-null", "B-null", "C-null", "D-null"))); + assertThat(results, equalTo(asList("A-null", "B-null", "C-null", "D-null", "E-null", "F-null"))); } @Test @@ -160,9 +185,14 @@ public void null_is_passed_as_context_to_map_loader_if_you_do_nothing() { loader.load("B"); loader.loadMany(asList("C", "D")); + Map keysAndContexts = new LinkedHashMap<>(); + keysAndContexts.put("E", null); + keysAndContexts.put("F", null); + loader.loadMany(keysAndContexts); + List results = loader.dispatchAndJoin(); - assertThat(results, equalTo(asList("A-null", "B-null", "C-null", "D-null"))); + assertThat(results, equalTo(asList("A-null", "B-null", "C-null", "D-null", "E-null", "F-null"))); } @Test diff --git a/src/test/java/org/dataloader/DataLoaderStatsTest.java b/src/test/java/org/dataloader/DataLoaderStatsTest.java index 06e8ae6..b8393e6 100644 --- a/src/test/java/org/dataloader/DataLoaderStatsTest.java +++ b/src/test/java/org/dataloader/DataLoaderStatsTest.java @@ -13,6 +13,7 @@ import java.util.ArrayList; import java.util.List; +import java.util.Map; import java.util.concurrent.CompletableFuture; import static java.util.Arrays.asList; @@ -118,9 +119,10 @@ public void stats_are_collected_with_caching_disabled() { loader.load("A"); loader.load("B"); loader.loadMany(asList("C", "D")); + loader.loadMany(Map.of("E", "E", "F", "F")); Statistics stats = loader.getStatistics(); - assertThat(stats.getLoadCount(), equalTo(4L)); + assertThat(stats.getLoadCount(), equalTo(6L)); assertThat(stats.getBatchInvokeCount(), equalTo(0L)); assertThat(stats.getBatchLoadCount(), equalTo(0L)); assertThat(stats.getCacheHitCount(), equalTo(0L)); @@ -128,9 +130,9 @@ public void stats_are_collected_with_caching_disabled() { loader.dispatch(); stats = loader.getStatistics(); - assertThat(stats.getLoadCount(), equalTo(4L)); + assertThat(stats.getLoadCount(), equalTo(6L)); assertThat(stats.getBatchInvokeCount(), equalTo(1L)); - assertThat(stats.getBatchLoadCount(), equalTo(4L)); + assertThat(stats.getBatchLoadCount(), equalTo(6L)); assertThat(stats.getCacheHitCount(), equalTo(0L)); loader.load("A"); @@ -139,9 +141,9 @@ public void stats_are_collected_with_caching_disabled() { loader.dispatch(); stats = loader.getStatistics(); - assertThat(stats.getLoadCount(), equalTo(6L)); + assertThat(stats.getLoadCount(), equalTo(8L)); assertThat(stats.getBatchInvokeCount(), equalTo(2L)); - assertThat(stats.getBatchLoadCount(), equalTo(6L)); + assertThat(stats.getBatchLoadCount(), equalTo(8L)); assertThat(stats.getCacheHitCount(), equalTo(0L)); } diff --git a/src/test/java/org/dataloader/DataLoaderTest.java b/src/test/java/org/dataloader/DataLoaderTest.java index 1ce34ea..c4ae883 100644 --- a/src/test/java/org/dataloader/DataLoaderTest.java +++ b/src/test/java/org/dataloader/DataLoaderTest.java @@ -35,24 +35,19 @@ import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; -import java.util.ArrayList; -import java.util.Collection; -import java.util.HashMap; -import java.util.List; -import java.util.Map; -import java.util.Optional; +import java.util.*; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionStage; import java.util.concurrent.ExecutionException; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Function; 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 java.util.Collections.*; import static java.util.concurrent.CompletableFuture.*; import static org.awaitility.Awaitility.await; import static org.dataloader.DataLoaderFactory.newDataLoader; @@ -66,10 +61,7 @@ import static org.dataloader.fixtures.TestKit.listFrom; import static org.dataloader.impl.CompletableFutureKit.cause; import static org.hamcrest.MatcherAssert.assertThat; -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.hamcrest.Matchers.*; import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.fail; @@ -116,19 +108,25 @@ public void basic_map_batch_loading() { }; DataLoader loader = DataLoaderFactory.newMappedDataLoader(evensOnlyMappedBatchLoader); + final List keys = asList("C", "D"); + final Map keysAndContexts = new LinkedHashMap<>(); + keysAndContexts.put("E", null); + keysAndContexts.put("F", null); + loader.load("A"); loader.load("B"); - loader.loadMany(asList("C", "D")); + loader.loadMany(keys); + loader.loadMany(keysAndContexts); List results = loader.dispatchAndJoin(); - assertThat(results.size(), equalTo(4)); - assertThat(results, equalTo(asList("A", null, "C", null))); + assertThat(results.size(), equalTo(6)); + assertThat(results, equalTo(asList("A", null, "C", null, "E", null))); } @ParameterizedTest @MethodSource("dataLoaderFactories") - public void should_Support_loading_multiple_keys_in_one_call(TestDataLoaderFactory factory) { + public void should_Support_loading_multiple_keys_in_one_call_via_list(TestDataLoaderFactory factory) { AtomicBoolean success = new AtomicBoolean(); DataLoader identityLoader = factory.idLoader(new DataLoaderOptions(), new ArrayList<>()); @@ -142,6 +140,26 @@ public void should_Support_loading_multiple_keys_in_one_call(TestDataLoaderFacto assertThat(futureAll.toCompletableFuture().join(), equalTo(asList(1, 2))); } + @ParameterizedTest + @MethodSource("dataLoaderFactories") + public void should_Support_loading_multiple_keys_in_one_call_via_map(TestDataLoaderFactory factory) { + AtomicBoolean success = new AtomicBoolean(); + DataLoader identityLoader = factory.idLoader(new DataLoaderOptions(), new ArrayList<>()); + + final Map keysAndContexts = new LinkedHashMap<>(); + keysAndContexts.put(1, null); + keysAndContexts.put(2, null); + + CompletionStage> futureAll = identityLoader.loadMany(keysAndContexts); + futureAll.thenAccept(promisedValues -> { + assertThat(promisedValues.size(), is(2)); + success.set(true); + }); + identityLoader.dispatch(); + await().untilAtomic(success, is(true)); + assertThat(futureAll.toCompletableFuture().join(), equalTo(Map.of(1, 1, 2, 2))); + } + @ParameterizedTest @MethodSource("dataLoaderFactories") public void should_Resolve_to_empty_list_when_no_keys_supplied(TestDataLoaderFactory factory) { @@ -159,7 +177,22 @@ public void should_Resolve_to_empty_list_when_no_keys_supplied(TestDataLoaderFac @ParameterizedTest @MethodSource("dataLoaderFactories") - public void should_Return_zero_entries_dispatched_when_no_keys_supplied(TestDataLoaderFactory factory) { + public void should_Resolve_to_empty_map_when_no_keys_supplied(TestDataLoaderFactory factory) { + AtomicBoolean success = new AtomicBoolean(); + DataLoader identityLoader = factory.idLoader(new DataLoaderOptions(), new ArrayList<>()); + CompletableFuture> futureEmpty = identityLoader.loadMany(emptyMap()); + futureEmpty.thenAccept(promisedValues -> { + assertThat(promisedValues.size(), is(0)); + success.set(true); + }); + identityLoader.dispatch(); + await().untilAtomic(success, is(true)); + assertThat(futureEmpty.join(), anEmptyMap()); + } + + @ParameterizedTest + @MethodSource("dataLoaderFactories") + public void should_Return_zero_entries_dispatched_when_no_keys_supplied_via_list(TestDataLoaderFactory factory) { AtomicBoolean success = new AtomicBoolean(); DataLoader identityLoader = factory.idLoader(new DataLoaderOptions(), new ArrayList<>()); CompletableFuture> futureEmpty = identityLoader.loadMany(emptyList()); @@ -172,6 +205,21 @@ public void should_Return_zero_entries_dispatched_when_no_keys_supplied(TestData assertThat(dispatchResult.getKeysCount(), equalTo(0)); } + @ParameterizedTest + @MethodSource("dataLoaderFactories") + public void should_Return_zero_entries_dispatched_when_no_keys_supplied_via_map(TestDataLoaderFactory factory) { + AtomicBoolean success = new AtomicBoolean(); + DataLoader identityLoader = factory.idLoader(new DataLoaderOptions(), new ArrayList<>()); + CompletableFuture> futureEmpty = identityLoader.loadMany(emptyMap()); + 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)); + } + @ParameterizedTest @MethodSource("dataLoaderFactories") public void should_Batch_multiple_requests(TestDataLoaderFactory factory) throws ExecutionException, InterruptedException { @@ -286,10 +334,17 @@ public void should_Cache_on_redispatch(TestDataLoaderFactory factory) throws Exe CompletableFuture> future2 = identityLoader.loadMany(asList("A", "B")); identityLoader.dispatch(); - await().until(() -> future1.isDone() && future2.isDone()); + Map keysAndContexts = new LinkedHashMap<>(); + keysAndContexts.put("A", null); + keysAndContexts.put("C", null); + CompletableFuture> future3 = identityLoader.loadMany(keysAndContexts); + identityLoader.dispatch(); + + await().until(() -> future1.isDone() && future2.isDone() && future3.isDone()); assertThat(future1.get(), equalTo("A")); assertThat(future2.get(), equalTo(asList("A", "B"))); - assertThat(loadCalls, equalTo(asList(singletonList("A"), singletonList("B")))); + assertThat(future3.get(), equalTo(keysAndContexts.entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getKey)))); + assertThat(loadCalls, equalTo(asList(singletonList("A"), singletonList("B"), singletonList("C")))); } @ParameterizedTest