Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit bce4817

Browse files
andimarekandimarek
authored
Merge pull request #4136 from graphql-java/chained-dl-perf
More Chained DataLoaders performance improvements
2 parents 676dae8 + 7598c55 commit bce4817

File tree

8 files changed

+120
-87
lines changed

8 files changed

+120
-87
lines changed

‎src/jmh/java/performance/DataLoaderPerformance.java‎

Lines changed: 10 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@
3535

3636
@State(Scope.Benchmark)
3737
@Warmup(iterations = 2, time = 5)
38-
@Measurement(iterations = 3)
38+
@Measurement(iterations = 5)
3939
@Fork(2)
4040
public class DataLoaderPerformance {
4141

@@ -481,25 +481,17 @@ public Pet(String id, String name, String ownerId, List<String> friendsIds) {
481481

482482

483483
static BatchLoader<String, Owner> ownerBatchLoader = list -> {
484+
// System.out.println("OwnerBatchLoader with " + list.size() );
484485
List<Owner> collect = list.stream().map(key -> {
485486
Owner owner = owners.get(key);
486-
try {
487-
Thread.sleep(50);
488-
} catch (InterruptedException e) {
489-
throw new RuntimeException(e);
490-
}
491487
return owner;
492488
}).collect(Collectors.toList());
493489
return CompletableFuture.completedFuture(collect);
494490
};
495491
static BatchLoader<String, Pet> petBatchLoader = list -> {
492+
// System.out.println("PetBatchLoader with list: " + list.size());
496493
List<Pet> collect = list.stream().map(key -> {
497494
Pet owner = pets.get(key);
498-
try {
499-
Thread.sleep(5);
500-
} catch (InterruptedException e) {
501-
throw new RuntimeException(e);
502-
}
503495
return owner;
504496
}).collect(Collectors.toList());
505497
return CompletableFuture.completedFuture(collect);
@@ -520,9 +512,6 @@ public void setup() {
520512
try {
521513
String sdl = PerformanceTestingUtils.loadResource("dataLoaderPerformanceSchema.graphqls");
522514

523-
524-
DataLoaderRegistry registry = new DataLoaderRegistry();
525-
526515
DataFetcher ownersDF = (env -> {
527516
// Load all 103 owners (O-1 through O-103)
528517
List<Object> allOwnerIds = List.of(
@@ -542,20 +531,20 @@ public void setup() {
542531
});
543532
DataFetcher petsDf = (env -> {
544533
Owner owner = env.getSource();
545-
return env.getDataLoader(petDLName).loadMany((List) owner.petIds);
546-
// .thenCompose((result) -> CompletableFuture.supplyAsync(() -> null).thenApply((__) -> result));
534+
return env.getDataLoader(petDLName).loadMany((List) owner.petIds)
535+
.thenCompose((result) -> CompletableFuture.supplyAsync(() -> null).thenApply((__) -> result));
547536
});
548537

549538
DataFetcher petFriendsDF = (env -> {
550539
Pet pet = env.getSource();
551-
return env.getDataLoader(petDLName).loadMany((List) pet.friendsIds);
552-
// .thenCompose((result) -> CompletableFuture.supplyAsync(() -> null).thenApply((__) -> result));
540+
return env.getDataLoader(petDLName).loadMany((List) pet.friendsIds)
541+
.thenCompose((result) -> CompletableFuture.supplyAsync(() -> null).thenApply((__) -> result));
553542
});
554543

555544
DataFetcher petOwnerDF = (env -> {
556545
Pet pet = env.getSource();
557-
return env.getDataLoader(ownerDLName).load(pet.ownerId);
558-
// .thenCompose((result) -> CompletableFuture.supplyAsync(() -> null).thenApply((__) -> result));
546+
return env.getDataLoader(ownerDLName).load(pet.ownerId)
547+
.thenCompose((result) -> CompletableFuture.supplyAsync(() -> null).thenApply((__) -> result));
559548
});
560549

561550

@@ -598,6 +587,7 @@ public void executeRequestWithDataLoaders(MyState myState, Blackhole blackhole)
598587
// .profileExecution(true)
599588
.build();
600589
executionInput.getGraphQLContext().put(DataLoaderDispatchingContextKeys.ENABLE_DATA_LOADER_CHAINING, true);
590+
// executionInput.getGraphQLContext().put(DataLoaderDispatchingContextKeys.ENABLE_DATA_LOADER_EXHAUSTED_DISPATCHING, true);
601591
ExecutionResult execute = myState.graphQL.execute(executionInput);
602592
// ProfilerResult profilerResult = executionInput.getGraphQLContext().get(ProfilerResult.PROFILER_CONTEXT_KEY);
603593
// System.out.println("execute: " + execute);

‎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/execution/ExecutionStrategy.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -448,6 +448,7 @@ private Object fetchField(GraphQLFieldDefinition fieldDef, ExecutionContext exec
448448
.selectionSet(fieldCollector)
449449
.queryDirectives(queryDirectives)
450450
.deferredCallContext(parameters.getDeferredCallContext())
451+
.level(parameters.getPath().getLevel())
451452
.build();
452453
});
453454

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,8 +44,8 @@ public static ResultPath rootPath() {
4444
private ResultPath() {
4545
parent = null;
4646
segment = null;
47-
this.toStringValue = initString();
4847
this.level = 0;
48+
this.toStringValue = initString();
4949
}
5050

5151
private ResultPath(ResultPath parent, String segment) {

‎src/main/java/graphql/execution/instrumentation/dataloader/PerLevelDataLoaderDispatchStrategy.java‎

Lines changed: 11 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -46,18 +46,18 @@ private static class ChainedDLStack {
4646
// a state for level points to a previous one
4747
// all the invocations that are linked together are the relevant invocations for the next dispatch
4848
private static class StateForLevel {
49-
final @Nullable DataLoaderInvocation dataLoaderInvocation;
49+
final @Nullable DataLoader dataLoader;
5050
final boolean dispatchingStarted;
5151
final boolean dispatchingFinished;
5252
final boolean currentlyDelayedDispatching;
5353
final @Nullable StateForLevel prev;
5454

55-
public StateForLevel(@Nullable DataLoaderInvocation dataLoaderInvocation,
55+
public StateForLevel(@Nullable DataLoader dataLoader,
5656
boolean dispatchingStarted,
5757
boolean dispatchingFinished,
5858
boolean currentlyDelayedDispatching,
5959
@Nullable StateForLevel prev) {
60-
this.dataLoaderInvocation = dataLoaderInvocation;
60+
this.dataLoader = dataLoader;
6161
this.dispatchingStarted = dispatchingStarted;
6262
this.dispatchingFinished = dispatchingFinished;
6363
this.currentlyDelayedDispatching = currentlyDelayedDispatching;
@@ -91,7 +91,7 @@ public StateForLevel(@Nullable DataLoaderInvocation dataLoaderInvocation,
9191
}
9292
}
9393

94-
if (currentState == null || currentState.dataLoaderInvocation == null) {
94+
if (currentState == null || currentState.dataLoader == null) {
9595
if (normalDispatchOrDelayed) {
9696
dispatchingFinished = true;
9797
} else {
@@ -108,8 +108,7 @@ public StateForLevel(@Nullable DataLoaderInvocation dataLoaderInvocation,
108108
}
109109

110110

111-
public boolean newDataLoaderInvocation(DataLoaderInvocation dataLoaderInvocation) {
112-
int level = dataLoaderInvocation.level;
111+
public boolean newDataLoaderInvocation(int level, DataLoader dataLoader) {
113112
AtomicReference<@Nullable StateForLevel> currentStateRef = stateMapPerLevel.computeIfAbsent(level, __ -> new AtomicReference<>());
114113
while (true) {
115114
StateForLevel currentState = currentStateRef.get();
@@ -132,7 +131,7 @@ public boolean newDataLoaderInvocation(DataLoaderInvocation dataLoaderInvocation
132131
currentlyDelayedDispatching = true;
133132
}
134133

135-
StateForLevel newState = new StateForLevel(dataLoaderInvocation, dispatchingStarted, dispatchingFinished, currentlyDelayedDispatching, currentState);
134+
StateForLevel newState = new StateForLevel(dataLoader, dispatchingStarted, dispatchingFinished, currentlyDelayedDispatching, currentState);
136135

137136
if (currentStateRef.compareAndSet(currentState, newState)) {
138137
return newDelayedInvocation;
@@ -487,20 +486,14 @@ private void dispatchAll(DataLoaderRegistry dataLoaderRegistry, int level) {
487486
private void dispatchDLCFImpl(Integer level, CallStack callStack, boolean normalOrDelayed, boolean chained) {
488487

489488
ChainedDLStack.StateForLevel stateForLevel = callStack.chainedDLStack.aboutToStartDispatching(level, normalOrDelayed, chained);
490-
if (stateForLevel == null || stateForLevel.dataLoaderInvocation == null) {
489+
if (stateForLevel == null || stateForLevel.dataLoader == null) {
491490
return;
492491
}
493492

494493
List<CompletableFuture> allDispatchedCFs = new ArrayList<>();
495-
while (stateForLevel != null && stateForLevel.dataLoaderInvocation != null) {
496-
final DataLoaderInvocation invocation = stateForLevel.dataLoaderInvocation;
497-
CompletableFuture<List> dispatch = invocation.dataLoader.dispatch();
494+
while (stateForLevel != null && stateForLevel.dataLoader != null) {
495+
CompletableFuture<List> dispatch = stateForLevel.dataLoader.dispatch();
498496
allDispatchedCFs.add(dispatch);
499-
dispatch.whenComplete((objects, throwable) -> {
500-
if (objects != null && objects.size() > 0) {
501-
profiler.batchLoadedNewStrategy(invocation.name, level, objects.size(), !normalOrDelayed, chained);
502-
}
503-
});
504497
stateForLevel = stateForLevel.prev;
505498
}
506499
CompletableFuture.allOf(allDispatchedCFs.toArray(new CompletableFuture[0]))
@@ -512,51 +505,19 @@ private void dispatchDLCFImpl(Integer level, CallStack callStack, boolean normal
512505
}
513506

514507

515-
public void newDataLoaderInvocation(String resultPath,
516-
int level,
508+
public void newDataLoaderInvocation(int level,
517509
DataLoader dataLoader,
518-
String dataLoaderName,
519-
Object key,
520510
@Nullable AlternativeCallContext alternativeCallContext) {
521511
if (!enableDataLoaderChaining) {
522512
return;
523513
}
524-
DataLoaderInvocation dataLoaderInvocation = new DataLoaderInvocation(resultPath, level, dataLoader, dataLoaderName, key);
525514
CallStack callStack = getCallStack(alternativeCallContext);
526-
boolean newDelayedInvocation = callStack.chainedDLStack.newDataLoaderInvocation(dataLoaderInvocation);
515+
boolean newDelayedInvocation = callStack.chainedDLStack.newDataLoaderInvocation(level, dataLoader);
527516
if (newDelayedInvocation) {
528517
dispatchDLCFImpl(level, callStack, false, false);
529518
}
530519
}
531520

532-
/**
533-
* A single data loader invocation.
534-
*/
535-
private static class DataLoaderInvocation {
536-
final String resultPath;
537-
final int level;
538-
final DataLoader dataLoader;
539-
final String name;
540-
final Object key;
541-
542-
public DataLoaderInvocation(String resultPath, int level, DataLoader dataLoader, String name, Object key) {
543-
this.resultPath = resultPath;
544-
this.level = level;
545-
this.dataLoader = dataLoader;
546-
this.name = name;
547-
this.key = key;
548-
}
549-
550-
@Override
551-
public String toString() {
552-
return "ResultPathWithDataLoader{" +
553-
"resultPath='" + resultPath + '\'' +
554-
", level=" + level +
555-
", key=" + key +
556-
", name='" + name + '\'' +
557-
'}';
558-
}
559-
}
560521

561522
}
562523

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

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,7 @@ public class DataFetchingEnvironmentImpl implements DataFetchingEnvironment {
6060
private final Document document;
6161
private final ImmutableMapWithNullValues<String, Object> variables;
6262
private final QueryDirectives queryDirectives;
63+
private final int level;
6364

6465
// used for internal() method
6566
private final DFEInternalState dfeInternalState;
@@ -86,6 +87,7 @@ private DataFetchingEnvironmentImpl(Builder builder) {
8687
this.document = builder.document;
8788
this.variables = builder.variables == null ? ImmutableMapWithNullValues.emptyMap() : builder.variables;
8889
this.queryDirectives = builder.queryDirectives;
90+
this.level = builder.level;
8991

9092
// internal state
9193
this.dfeInternalState = new DFEInternalState(builder.dataLoaderDispatchStrategy, builder.alternativeCallContext, builder.profiler);
@@ -278,6 +280,10 @@ public String toString() {
278280
'}';
279281
}
280282

283+
public int getLevel() {
284+
return level;
285+
}
286+
281287
@NullUnmarked
282288
public static class Builder {
283289

@@ -305,6 +311,7 @@ public static class Builder {
305311
private DataLoaderDispatchStrategy dataLoaderDispatchStrategy;
306312
private Profiler profiler;
307313
private AlternativeCallContext alternativeCallContext;
314+
private int level;
308315

309316
public Builder(DataFetchingEnvironmentImpl env) {
310317
this.source = env.source;
@@ -331,6 +338,7 @@ public Builder(DataFetchingEnvironmentImpl env) {
331338
this.dataLoaderDispatchStrategy = env.dfeInternalState.dataLoaderDispatchStrategy;
332339
this.profiler = env.dfeInternalState.profiler;
333340
this.alternativeCallContext = env.dfeInternalState.alternativeCallContext;
341+
this.level = env.level;
334342
}
335343

336344
public Builder() {
@@ -468,6 +476,11 @@ public Builder profiler(Profiler profiler) {
468476
this.profiler = profiler;
469477
return this;
470478
}
479+
480+
public Builder level(int level) {
481+
this.level = level;
482+
return this;
483+
}
471484
}
472485

473486
@Internal

0 commit comments

Comments
 (0)