Skip to content

make dataloader work inside defer blocks #3980

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
andimarek merged 15 commits into from
May 26, 2025
Merged

make dataloader work inside defer blocks #3980

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
15 commits
Select commit Hold shift + click to select a range
15f6952
make dataloader work inside defer blocks
andimarek May 22, 2025
da6886a
green tests
andimarek May 22, 2025
018248a
update gradle wrapper validation
andimarek May 22, 2025
9aaa118
update gradle wrapper validation
andimarek May 22, 2025
e7e4c36
Merge branch 'master' into defer-dataloader
andimarek May 22, 2025
386ce8f
cleanup
andimarek May 22, 2025
fcf7261
more complex test
andimarek May 23, 2025
17911d3
cleanup
andimarek May 23, 2025
5c64ad4
cleanup and docs
andimarek May 23, 2025
c8f4362
cleanup
andimarek May 23, 2025
e0181e3
testing
andimarek May 23, 2025
f345f14
cleanup
andimarek May 23, 2025
c12be77
cleanup
andimarek May 23, 2025
aa206d7
cleanup
andimarek May 23, 2025
a7cc794
PR feedback
andimarek May 26, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
green tests
  • Loading branch information
@andimarek
andimarek committed May 22, 2025
commit da6886a15022cdf549956f7b1ed095a92e39e447
4 changes: 3 additions & 1 deletion src/main/java/graphql/execution/AsyncExecutionStrategy.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,6 @@ public AsyncExecutionStrategy(DataFetcherExceptionHandler exceptionHandler) {
@SuppressWarnings("FutureReturnValueIgnored")
public CompletableFuture<ExecutionResult> execute(ExecutionContext executionContext, ExecutionStrategyParameters parameters) throws NonNullableFieldWasNullException {
DataLoaderDispatchStrategy dataLoaderDispatcherStrategy = executionContext.getDataLoaderDispatcherStrategy();
dataLoaderDispatcherStrategy.executionStrategy(executionContext, parameters);
Instrumentation instrumentation = executionContext.getInstrumentation();
InstrumentationExecutionStrategyParameters instrumentationParameters = new InstrumentationExecutionStrategyParameters(executionContext, parameters);

Expand All @@ -54,6 +53,9 @@ public CompletableFuture<ExecutionResult> execute(ExecutionContext executionCont
}

DeferredExecutionSupport deferredExecutionSupport = createDeferredExecutionSupport(executionContext, parameters);

dataLoaderDispatcherStrategy.executionStrategy(executionContext, parameters, deferredExecutionSupport.getNonDeferredFieldNames(fieldNames).size());

Async.CombinedBuilder<FieldValueInfo> futures = getAsyncFieldValueInfo(executionContext, parameters, deferredExecutionSupport);

CompletableFuture<ExecutionResult> overallResult = new CompletableFuture<>();
Expand Down
4 changes: 2 additions & 2 deletions src/main/java/graphql/execution/DataLoaderDispatchStrategy.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ public interface DataLoaderDispatchStrategy {
};


default void executionStrategy(ExecutionContext executionContext, ExecutionStrategyParameters parameters) {
default void executionStrategy(ExecutionContext executionContext, ExecutionStrategyParameters parameters, int fieldCount) {

}

Expand All @@ -31,7 +31,7 @@ default void executionStrategyOnFieldValuesException(Throwable t, ExecutionStrat
}


default void executeObject(ExecutionContext executionContext, ExecutionStrategyParameters executionStrategyParameters) {
default void executeObject(ExecutionContext executionContext, ExecutionStrategyParameters executionStrategyParameters, int fieldCount) {

}

Expand Down
2 changes: 1 addition & 1 deletion src/main/java/graphql/execution/ExecutionStrategy.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -195,7 +195,6 @@ public static String mkNameForPath(List<Field> currentField) {
@DuckTyped(shape = "CompletableFuture<Map<String, Object>> | Map<String, Object>")
protected Object executeObject(ExecutionContext executionContext, ExecutionStrategyParameters parameters) throws NonNullableFieldWasNullException {
DataLoaderDispatchStrategy dataLoaderDispatcherStrategy = executionContext.getDataLoaderDispatcherStrategy();
dataLoaderDispatcherStrategy.executeObject(executionContext, parameters);
Instrumentation instrumentation = executionContext.getInstrumentation();
InstrumentationExecutionStrategyParameters instrumentationParameters = new InstrumentationExecutionStrategyParameters(executionContext, parameters);

Expand All @@ -210,6 +209,7 @@ protected Object executeObject(ExecutionContext executionContext, ExecutionStrat

CompletableFuture<Map<String, Object>> overallResult = new CompletableFuture<>();
List<String> fieldsExecutedOnInitialResult = deferredExecutionSupport.getNonDeferredFieldNames(fieldNames);
dataLoaderDispatcherStrategy.executeObject(executionContext, parameters, fieldsExecutedOnInitialResult.size());
BiConsumer<List<Object>, Throwable> handleResultsConsumer = buildFieldValueMap(fieldsExecutedOnInitialResult, overallResult, executionContext);

resolveObjectCtx.onDispatched();
Expand Down
7 changes: 7 additions & 0 deletions src/main/java/graphql/execution/incremental/DeferredCallContext.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

import graphql.GraphQLError;
import graphql.Internal;
import graphql.VisibleForTesting;

import java.util.List;
import java.util.concurrent.CopyOnWriteArrayList;
Expand All @@ -26,6 +27,12 @@ public DeferredCallContext(int startLevel, int fields) {
this.fields = fields;
}

@VisibleForTesting
public DeferredCallContext() {
this.startLevel = 0;
this.fields = 0;
}

public int getStartLevel() {
return startLevel;
}
Expand Down
1 change: 0 additions & 1 deletion src/main/java/graphql/execution/incremental/DeferredExecutionSupport.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -117,7 +117,6 @@ public Set<IncrementalCall<? extends IncrementalPayload>> createCalls() {

private DeferredFragmentCall createDeferredFragmentCall(DeferredExecution deferredExecution) {
int level = parameters.getPath().getLevel() + 1;
System.out.println("new DeferredFragmentCall for level " + level + " with fields " + deferredFields.size());
DeferredCallContext deferredCallContext = new DeferredCallContext(level, deferredFields.size());

List<MergedField> mergedFields = deferredExecutionToFields.get(deferredExecution);
Expand Down
85 changes: 48 additions & 37 deletions ...java/graphql/execution/instrumentation/dataloader/PerLevelDataLoaderDispatchStrategy.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@
@NullMarked
public class PerLevelDataLoaderDispatchStrategy implements DataLoaderDispatchStrategy {

private final CallStack callStack;
private final CallStack initialCallStack;
private final ExecutionContext executionContext;
private final long batchWindowNs;
private final boolean enableDataLoaderChaining;
Expand Down Expand Up @@ -90,12 +90,12 @@ private static class CallStack {
private boolean batchWindowOpen;


private List<FieldValueInfo> deferredFragmentRootFieldsFetched;
private final List<FieldValueInfo> deferredFragmentRootFieldsFetched = new ArrayList<>();

public CallStack() {
// in the first level there is only one sub selection,
// so we only expect one execute object call (which is actually an executionStrategy call)
expectedExecuteObjectCallsPerLevel.set(1, 1);
expectedExecuteObjectCallsPerLevel.set(0, 1);
}

public void addResultPathWithDataLoader(int level, ResultPathWithDataLoader resultPathWithDataLoader) {
Expand Down Expand Up @@ -148,8 +148,8 @@ boolean allExecuteObjectCallsHappened(int level) {
return happenedExecuteObjectCallsPerLevel.get(level) == expectedExecuteObjectCallsPerLevel.get(level);
}

boolean allSubSelectionsFetchingHappened(int level) {
return happenedOnFieldValueCallsPerLevel.get(level) == expectedExecuteObjectCallsPerLevel.get(level);
boolean allSubSelectionsFetchingHappened(int subSelectionLevel) {
return happenedOnFieldValueCallsPerLevel.get(subSelectionLevel) == expectedExecuteObjectCallsPerLevel.get(subSelectionLevel - 1);
}

boolean allFetchesHappened(int level) {
Expand Down Expand Up @@ -181,7 +181,7 @@ public void setDispatchedLevel(int level) {
}

public PerLevelDataLoaderDispatchStrategy(ExecutionContext executionContext) {
this.callStack = new CallStack();
this.initialCallStack = new CallStack();
this.executionContext = executionContext;

GraphQLContext graphQLContext = executionContext.getGraphQLContext();
Expand All @@ -204,21 +204,22 @@ public void executeDeferredOnFieldValueInfo(FieldValueInfo fieldValueInfo, Execu
}

@Override
public void executionStrategy(ExecutionContext executionContext, ExecutionStrategyParameters parameters) {
public void executionStrategy(ExecutionContext executionContext, ExecutionStrategyParameters parameters, int fieldCount) {
Assert.assertTrue(parameters.getExecutionStepInfo().getPath().isRootPath());
increaseHappenedExecuteObjectAndIncreaseExpectedFetchCount(1, parameters, callStack);
increaseHappenedExecuteObjectAndIncreaseExpectedFetchCount(0, fieldCount, initialCallStack);
}

@Override
public void executionSerialStrategy(ExecutionContext executionContext, ExecutionStrategyParameters parameters) {
CallStack callStack = getCallStack(parameters);
resetCallStack(callStack);
increaseHappenedExecuteObjectAndIncreaseExpectedFetchCount(1, 1, callStack);
increaseHappenedExecuteObjectAndIncreaseExpectedFetchCount(0, 1, callStack);
}

@Override
public void executionStrategyOnFieldValuesInfo(List<FieldValueInfo> fieldValueInfoList, ExecutionStrategyParameters parameters) {
CallStack callStack = getCallStack(parameters);
// the root fields are the root sub selection on level 1
onFieldValuesInfoDispatchIfNeeded(fieldValueInfoList, 1, callStack);
}

Expand All @@ -236,68 +237,72 @@ private CallStack getCallStack(ExecutionStrategyParameters parameters) {

private CallStack getCallStack(@Nullable DeferredCallContext deferredCallContext) {
if (deferredCallContext == null) {
return this.callStack;
return this.initialCallStack;
} else {
return callStackMap.computeIfAbsent(deferredCallContext, k -> {
CallStack callStack = new CallStack();
int startLevel = deferredCallContext.getStartLevel();
int fields = deferredCallContext.getFields();
callStack.lock.runLocked(() -> {
callStack.expectedExecuteObjectCallsPerLevel.set(0, 0); // set to 1 in the constructor of CallStack
callStack.expectedExecuteObjectCallsPerLevel.set(startLevel - 1, 1);
callStack.happenedExecuteObjectCallsPerLevel.set(startLevel - 1, 1);
callStack.highestReadyLevel = startLevel - 1;
callStack.increaseExpectedFetchCount(startLevel, fields);
// we make sure that startLevel-1 is considered done
callStack.expectedExecuteObjectCallsPerLevel.set(startLevel - 1, 0);
callStack.happenedOnFieldValueCallsPerLevel.set(startLevel - 1, 0);
callStack.highestReadyLevel = startLevel - 1;
});
return callStack;
});
}
}

@Override
public void executeObject(ExecutionContext executionContext, ExecutionStrategyParameters parameters) {
public void executeObject(ExecutionContext executionContext, ExecutionStrategyParameters parameters, int fieldCount) {
CallStack callStack = getCallStack(parameters);
int curLevel = parameters.getExecutionStepInfo().getPath().getLevel() + 1;
increaseHappenedExecuteObjectAndIncreaseExpectedFetchCount(curLevel, parameters, callStack);
int curLevel = parameters.getPath().getLevel();
increaseHappenedExecuteObjectAndIncreaseExpectedFetchCount(curLevel, fieldCount, callStack);
}

@Override
public void executeObjectOnFieldValuesInfo(List<FieldValueInfo> fieldValueInfoList, ExecutionStrategyParameters parameters) {
public void executeObjectOnFieldValuesInfo
(List<FieldValueInfo> fieldValueInfoList, ExecutionStrategyParameters parameters) {
// the level of the sub selection that is fully fetched is one level more than parameters level
int curLevel = parameters.getPath().getLevel() + 1;
CallStack callStack = getCallStack(parameters);
onFieldValuesInfoDispatchIfNeeded(fieldValueInfoList, curLevel, callStack);
}

@Override
public void deferredOnFieldValue(String resultKey, FieldValueInfo fieldValueInfo, Throwable throwable, ExecutionStrategyParameters parameters) {
public void deferredOnFieldValue(String resultKey, FieldValueInfo fieldValueInfo, Throwable
throwable, ExecutionStrategyParameters parameters) {
CallStack callStack = getCallStack(parameters);
boolean ready = callStack.lock.callLocked(() -> {
callStack.deferredFragmentRootFieldsFetched.add(fieldValueInfo);
return callStack.deferredFragmentRootFieldsFetched.size() == parameters.getDeferredCallContext().getFields();
});
if (ready) {
int curLevel = parameters.getPath().getLevel() + 1;
int curLevel = parameters.getPath().getLevel();
onFieldValuesInfoDispatchIfNeeded(callStack.deferredFragmentRootFieldsFetched, curLevel, callStack);
}
}

@Override
public void executeObjectOnFieldValuesException(Throwable t, ExecutionStrategyParameters parameters) {
CallStack callStack = getCallStack(parameters);
// the level of the sub selection that is errored is one level more than parameters level
int curLevel = parameters.getPath().getLevel() + 1;
callStack.lock.runLocked(() ->
callStack.increaseHappenedOnFieldValueCalls(curLevel)
);
}

private void increaseHappenedExecuteObjectAndIncreaseExpectedFetchCount(int curLevel, ExecutionStrategyParameters executionStrategyParameters, CallStack callStack) {
increaseHappenedExecuteObjectAndIncreaseExpectedFetchCount(curLevel, executionStrategyParameters.getFields().size(), callStack);
}

private void increaseHappenedExecuteObjectAndIncreaseExpectedFetchCount(int curLevel, int fieldCount, CallStack callStack) {
private void increaseHappenedExecuteObjectAndIncreaseExpectedFetchCount(int curLevel,
int fieldCount,
CallStack callStack) {
callStack.lock.runLocked(() -> {
callStack.increaseHappenedExecuteObjectCalls(curLevel);
callStack.increaseExpectedFetchCount(curLevel, fieldCount);
callStack.increaseExpectedFetchCount(curLevel + 1, fieldCount);
});
}

Expand All @@ -309,7 +314,7 @@ private void resetCallStack(CallStack callStack) {
callStack.clearFetchCount();
callStack.clearHappenedExecuteObjectCalls();
callStack.clearHappenedOnFieldValueCalls();
callStack.expectedExecuteObjectCallsPerLevel.set(1, 1);
callStack.expectedExecuteObjectCallsPerLevel.set(0, 1);
callStack.dispatchingFinishedPerLevel.clear();
callStack.dispatchingStartedPerLevel.clear();
callStack.allResultPathWithDataLoader.clear();
Expand All @@ -320,9 +325,11 @@ private void resetCallStack(CallStack callStack) {
});
}

private void onFieldValuesInfoDispatchIfNeeded(List<FieldValueInfo> fieldValueInfoList, int curLevel, CallStack callStack) {
private void onFieldValuesInfoDispatchIfNeeded(List<FieldValueInfo> fieldValueInfoList,
int subSelectionLevel,
CallStack callStack) {
Integer dispatchLevel = callStack.lock.callLocked(() ->
handleOnFieldValuesInfo(fieldValueInfoList, curLevel, callStack)
handleSubSelectionFetched(fieldValueInfoList, subSelectionLevel, callStack)
);
// the handle on field values check for the next level if it is ready
if (dispatchLevel != null) {
Expand All @@ -333,11 +340,12 @@ private void onFieldValuesInfoDispatchIfNeeded(List<FieldValueInfo> fieldValueIn
//
// thread safety: called with callStack.lock
//
private Integer handleOnFieldValuesInfo(List<FieldValueInfo> fieldValueInfos, int curLevel, CallStack callStack) {
callStack.increaseHappenedOnFieldValueCalls(curLevel);
private Integer handleSubSelectionFetched(List<FieldValueInfo> fieldValueInfos, int subSelectionLevel, CallStack
callStack) {
callStack.increaseHappenedOnFieldValueCalls(subSelectionLevel);
int expectedOnObjectCalls = getObjectCountForList(fieldValueInfos);
// on the next level we expect the following on object calls because we found non null objects
callStack.increaseExpectedExecuteObjectCalls(curLevel + 1, expectedOnObjectCalls);
// we expect on the level of the current sub selection #expectedOnObjectCalls execute object calls
callStack.increaseExpectedExecuteObjectCalls(subSelectionLevel, expectedOnObjectCalls);
// maybe the object calls happened already (because the DataFetcher return directly values synchronously)
// therefore we check the next levels if they are ready
// this means we could skip some level because the higher level is also already ready,
Expand All @@ -347,7 +355,7 @@ private Integer handleOnFieldValuesInfo(List<FieldValueInfo> fieldValueInfos, in
// if data loader chaining is disabled (the old algo) the level we dispatch is not really relevant as
// we dispatch the whole registry anyway

return getHighestReadyLevel(curLevel + 1, callStack);
return getHighestReadyLevel(subSelectionLevel + 1, callStack);
}

/**
Expand Down Expand Up @@ -431,14 +439,16 @@ private boolean checkLevelImpl(int level, CallStack callStack) {
if (callStack.expectedFetchCountPerLevel.get(level) == 0) {
return false;
}
// level 1 is special: there is no previous sub selections
// and the expected execution object calls is always 1
if (level > 1 && !callStack.allSubSelectionsFetchingHappened(level - 1)) {

// first we make sure that the expected fetch count is correct
// by verifying that the parent level all execute object + sub selection were fetched
if (!callStack.allExecuteObjectCallsHappened(level - 1)) {
return false;
}
if (!callStack.allExecuteObjectCallsHappened(level)) {
if (level > 1 && !callStack.allSubSelectionsFetchingHappened(level - 1)) {
return false;
}
// the main check: all fetches must have happened
if (!callStack.allFetchesHappened(level)) {
return false;
}
Expand Down Expand Up @@ -503,7 +513,8 @@ public void dispatchDLCFImpl(Set<String> resultPathsToDispatch, Integer level, C
}


public void newDataLoaderLoadCall(String resultPath, int level, DataLoader dataLoader, String dataLoaderName, Object key, @Nullable DeferredCallContext deferredCallContext) {
public void newDataLoaderLoadCall(String resultPath, int level, DataLoader dataLoader, String
dataLoaderName, Object key, @Nullable DeferredCallContext deferredCallContext) {
if (!enableDataLoaderChaining) {
return;
}
Expand Down
1 change: 1 addition & 0 deletions src/test/groovy/graphql/execution/incremental/IncrementalCallStateDeferTest.groovy
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@ class IncrementalCallStateDeferTest extends Specification {
results[2].incremental[0].data["a"] == "A"
}

// flaky
Comment thread
andimarek marked this conversation as resolved.
Outdated
def "calls within calls are enqueued correctly"() {
given:
def incrementalCallState = new IncrementalCallState()
Expand Down
2 changes: 2 additions & 0 deletions src/test/groovy/graphql/execution/instrumentation/dataloader/BatchCompareDataFetchers.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,7 @@ private static List<List<Department>> getDepartmentsForShops(List<Shop> shops) {


private BatchLoader<String, List<Department>> departmentsForShopsBatchLoader = ids -> maybeAsyncWithSleep(() -> {
System.out.println("departments for shops batch loader called with ids: " + ids);
Comment thread
andimarek marked this conversation as resolved.
Outdated
departmentsForShopsBatchLoaderCounter.incrementAndGet();
List<Shop> shopList = new ArrayList<>();
for (String id : ids) {
Expand Down Expand Up @@ -128,6 +129,7 @@ private static List<List<Product>> getProductsForDepartments(List<Department> de
}

private BatchLoader<String, List<Product>> productsForDepartmentsBatchLoader = ids -> maybeAsyncWithSleep(() -> {
System.out.println("products for deparments batch loader called with ids: " + ids);
productsForDepartmentsBatchLoaderCounter.incrementAndGet();
List<Department> d = ids.stream().map(departments::get).collect(Collectors.toList());
return completedFuture(getProductsForDepartments(d));
Expand Down
12 changes: 6 additions & 6 deletions ...test/groovy/graphql/execution/instrumentation/dataloader/DataLoaderPerformanceData.groovy
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -68,14 +68,14 @@ class DataLoaderPerformanceData {
static String getQuery(boolean deferDepartments, boolean deferProducts) {
return """
query {
shops {
id name
shops { # 1
id name # 2
... @defer(if: $deferDepartments) {
departments {
id name
departments { # 2
id name # 3
... @defer(if: $deferProducts) {
products {
id name
products { # 3
id name # 4
}
}
}
Expand Down
Loading