Skip to content

Commit 7598c55

Browse files
andimarekandimarek
committed
more performance
1 parent dd932bf commit 7598c55

File tree

2 files changed

+80
-2
lines changed

2 files changed

+80
-2
lines changed

‎src/main/java/graphql/execution/Async.java‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,14 +11,17 @@
1111
import java.util.Collections;
1212
import java.util.Iterator;
1313
import java.util.List;
14+
import java.util.Map;
1415
import java.util.concurrent.CompletableFuture;
1516
import java.util.concurrent.CompletionException;
1617
import java.util.concurrent.CompletionStage;
1718
import java.util.function.BiFunction;
1819
import java.util.function.Function;
1920
import java.util.function.Supplier;
21+
import java.util.stream.Collectors;
2022

2123
import static graphql.Assert.assertTrue;
24+
import static java.util.stream.Collectors.toList;
2225

2326
@Internal
2427
@SuppressWarnings("FutureReturnValueIgnored")
@@ -408,4 +411,24 @@ public static <T> CompletableFuture<T> exceptionallyCompletedFuture(Throwable ex
408411
public static <T> @NonNull CompletableFuture<T> orNullCompletedFuture(@Nullable CompletableFuture<T> completableFuture) {
409412
return completableFuture != null ? completableFuture : CompletableFuture.completedFuture(null);
410413
}
414+
415+
public static <T> CompletableFuture<List<T>> allOf(List<CompletableFuture<T>> cfs) {
416+
return CompletableFuture.allOf(cfs.toArray(CompletableFuture[]::new))
417+
.thenApply(v -> cfs.stream()
418+
.map(CompletableFuture::join)
419+
.collect(toList())
420+
);
421+
}
422+
423+
public static <K, V> CompletableFuture<Map<K, V>> allOf(Map<K, CompletableFuture<V>> cfs) {
424+
return CompletableFuture.allOf(cfs.values().toArray(CompletableFuture[]::new))
425+
.thenApply(v -> cfs.entrySet().stream()
426+
.collect(
427+
Collectors.toMap(
428+
Map.Entry::getKey,
429+
task -> task.getValue().join())
430+
)
431+
);
432+
}
433+
411434
}

‎src/main/java/graphql/schema/DataLoaderWithContext.java‎

Lines changed: 57 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package graphql.schema;
22

33
import graphql.Internal;
4+
import graphql.execution.Async;
45
import graphql.execution.incremental.AlternativeCallContext;
56
import graphql.execution.instrumentation.dataloader.ExhaustedDataLoaderDispatchStrategy;
67
import graphql.execution.instrumentation.dataloader.PerLevelDataLoaderDispatchStrategy;
@@ -10,8 +11,14 @@
1011
import org.jspecify.annotations.NullMarked;
1112
import org.jspecify.annotations.Nullable;
1213

14+
import java.util.ArrayList;
15+
import java.util.HashMap;
16+
import java.util.List;
17+
import java.util.Map;
1318
import java.util.concurrent.CompletableFuture;
1419

20+
import static graphql.Assert.assertNotNull;
21+
1522
@Internal
1623
@NullMarked
1724
public class DataLoaderWithContext<K, V> extends DelegatingDataLoader<K, V> {
@@ -29,9 +36,54 @@ public CompletableFuture<V> load(@NonNull K key, @Nullable Object keyContext) {
2936
// calling super.load() is important, because otherwise the data loader will sometimes called
3037
// later than the dispatch, which results in a hanging DL
3138
CompletableFuture<V> result = super.load(key, keyContext);
39+
newDataLoaderInvocation();
40+
return result;
41+
}
42+
43+
44+
@Override
45+
public CompletableFuture<List<V>> loadMany(List<K> keys, List<Object> keyContexts) {
46+
assertNotNull(keys);
47+
assertNotNull(keyContexts);
48+
49+
CompletableFuture<List<V>> result;
50+
synchronized (this) {
51+
List<CompletableFuture<V>> collect = new ArrayList<>(keys.size());
52+
for (int i = 0; i < keys.size(); i++) {
53+
K key = keys.get(i);
54+
Object keyContext = null;
55+
if (i < keyContexts.size()) {
56+
keyContext = keyContexts.get(i);
57+
}
58+
collect.add(delegate.load(key, keyContext));
59+
}
60+
result = Async.allOf(collect);
61+
}
62+
newDataLoaderInvocation();
63+
return result;
64+
}
65+
66+
@Override
67+
public CompletableFuture<Map<K, V>> loadMany(Map<K, ?> keysAndContexts) {
68+
assertNotNull(keysAndContexts);
69+
70+
CompletableFuture<Map<K, V>> result;
71+
synchronized (this) {
72+
Map<K, CompletableFuture<V>> collect = new HashMap<>(keysAndContexts.size());
73+
for (Map.Entry<K, ?> entry : keysAndContexts.entrySet()) {
74+
K key = entry.getKey();
75+
Object keyContext = entry.getValue();
76+
collect.put(key, delegate.load(key, keyContext));
77+
}
78+
result = Async.allOf(collect);
79+
}
80+
newDataLoaderInvocation();
81+
return result;
82+
}
83+
84+
private void newDataLoaderInvocation() {
3285
DataFetchingEnvironmentImpl dfeImpl = (DataFetchingEnvironmentImpl) dfe;
3386
DataFetchingEnvironmentImpl.DFEInternalState dfeInternalState = (DataFetchingEnvironmentImpl.DFEInternalState) dfeImpl.toInternal();
34-
dfeInternalState.getProfiler().dataLoaderUsed(dataLoaderName);
3587
if (dfeInternalState.getDataLoaderDispatchStrategy() instanceof PerLevelDataLoaderDispatchStrategy) {
3688
AlternativeCallContext alternativeCallContext = dfeInternalState.getDeferredCallContext();
3789
int level = dfeImpl.getLevel();
@@ -40,6 +92,9 @@ public CompletableFuture<V> load(@NonNull K key, @Nullable Object keyContext) {
4092
AlternativeCallContext alternativeCallContext = dfeInternalState.getDeferredCallContext();
4193
((ExhaustedDataLoaderDispatchStrategy) dfeInternalState.dataLoaderDispatchStrategy).newDataLoaderInvocation(alternativeCallContext);
4294
}
43-
return result;
4495
}
96+
97+
98+
99+
45100
}

0 commit comments

Comments
 (0)