Skip to content

Commit 914918d

Browse files
andimarekandimarek
committed
optimize chained DL code
1 parent a0181b8 commit 914918d

5 files changed

Lines changed: 32 additions & 62 deletions

File tree

‎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: 5 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -39,23 +39,27 @@ public static ResultPath rootPath() {
3939
// hash is effective immutable but lazily initialized similar to the hash code of java.lang.String
4040
private int hash;
4141
private final String toStringValue;
42+
private final int level;
4243

4344
private ResultPath() {
4445
parent = null;
4546
segment = null;
47+
this.level = 0;
4648
this.toStringValue = initString();
4749
}
4850

4951
private ResultPath(ResultPath parent, String segment) {
5052
this.parent = assertNotNull(parent, () -> "Must provide a parent path");
5153
this.segment = assertNotNull(segment, () -> "Must provide a sub path");
5254
this.toStringValue = initString();
55+
this.level = parent.level + 1;
5356
}
5457

5558
private ResultPath(ResultPath parent, int segment) {
5659
this.parent = assertNotNull(parent, () -> "Must provide a parent path");
5760
this.segment = segment;
5861
this.toStringValue = initString();
62+
this.level = parent.level;
5963
}
6064

6165
private String initString() {
@@ -71,15 +75,7 @@ private String initString() {
7175
}
7276

7377
public int getLevel() {
74-
int counter = 0;
75-
ResultPath currentPath = this;
76-
while (currentPath != null) {
77-
if (currentPath.segment instanceof String) {
78-
counter++;
79-
}
80-
currentPath = currentPath.parent;
81-
}
82-
return counter;
78+
return level;
8379
}
8480

8581
public ResultPath getPathWithoutListEnd() {

‎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

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

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -35,9 +35,8 @@ public CompletableFuture<V> load(@NonNull K key, @Nullable Object keyContext) {
3535
dfeInternalState.getProfiler().dataLoaderUsed(dataLoaderName);
3636
if (dfeInternalState.getDataLoaderDispatchStrategy() instanceof PerLevelDataLoaderDispatchStrategy) {
3737
AlternativeCallContext alternativeCallContext = dfeInternalState.getDeferredCallContext();
38-
int level = dfe.getExecutionStepInfo().getPath().getLevel();
39-
String path = dfe.getExecutionStepInfo().getPath().toString();
40-
((PerLevelDataLoaderDispatchStrategy) dfeInternalState.dataLoaderDispatchStrategy).newDataLoaderInvocation(path, level, delegate, dataLoaderName, key, alternativeCallContext);
38+
int level = dfeImpl.getLevel();
39+
((PerLevelDataLoaderDispatchStrategy) dfeInternalState.dataLoaderDispatchStrategy).newDataLoaderInvocation(level, delegate, alternativeCallContext);
4140
} else if (dfeInternalState.getDataLoaderDispatchStrategy() instanceof ExhaustedDataLoaderDispatchStrategy) {
4241
AlternativeCallContext alternativeCallContext = dfeInternalState.getDeferredCallContext();
4342
((ExhaustedDataLoaderDispatchStrategy) dfeInternalState.dataLoaderDispatchStrategy).newDataLoaderInvocation(alternativeCallContext);

0 commit comments

Comments
 (0)