Skip to content

JS like Batching #3873

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

Closed
andimarek wants to merge 3 commits into from
Closed

JS like Batching #3873

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
0a27e76
introducing of a special CompletableFuture enabling scheduling of nes…
andimarek Mar 30, 2025
f655893
replace countdown latch with an async solution
andimarek Mar 30, 2025
dc19532
only use window batching for DataLoaderCFs
andimarek Mar 30, 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
4 changes: 3 additions & 1 deletion src/main/java/graphql/execution/AbstractAsyncExecutionStrategy.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@
import graphql.ExecutionResultImpl;
import graphql.PublicSpi;

import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
Expand All @@ -24,8 +23,10 @@ public AbstractAsyncExecutionStrategy(DataFetcherExceptionHandler dataFetcherExc

protected BiConsumer<List<Object>, Throwable> handleResults(ExecutionContext executionContext, List<String> fieldNames, CompletableFuture<ExecutionResult> overallResult) {
return (List<Object> results, Throwable exception) -> {
executionContext.running();
if (exception != null) {
handleNonNullException(executionContext, overallResult, exception);
executionContext.finished();
return;
}
Map<String, Object> resolvedValuesByField = Maps.newLinkedHashMapWithExpectedSize(fieldNames.size());
Expand All @@ -35,6 +36,7 @@ protected BiConsumer<List<Object>, Throwable> handleResults(ExecutionContext exe
resolvedValuesByField.put(fieldName, result);
}
overallResult.complete(new ExecutionResultImpl(resolvedValuesByField, executionContext.getErrors()));
executionContext.finished();
};
}
}
8 changes: 8 additions & 0 deletions src/main/java/graphql/execution/AsyncExecutionStrategy.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@ public AsyncExecutionStrategy(DataFetcherExceptionHandler exceptionHandler) {
@Override
@SuppressWarnings("FutureReturnValueIgnored")
public CompletableFuture<ExecutionResult> execute(ExecutionContext executionContext, ExecutionStrategyParameters parameters) throws NonNullableFieldWasNullException {
executionContext.running();
DataLoaderDispatchStrategy dataLoaderDispatcherStrategy = executionContext.getDataLoaderDispatcherStrategy();
dataLoaderDispatcherStrategy.executionStrategy(executionContext, parameters);
Instrumentation instrumentation = executionContext.getInstrumentation();
Expand All @@ -50,6 +51,7 @@ public CompletableFuture<ExecutionResult> execute(ExecutionContext executionCont

Optional<ExecutionResult> isNotSensible = Introspection.isIntrospectionSensible(fields, executionContext);
if (isNotSensible.isPresent()) {
executionContext.finished();
return CompletableFuture.completedFuture(isNotSensible.get());
}

Expand All @@ -60,11 +62,13 @@ public CompletableFuture<ExecutionResult> execute(ExecutionContext executionCont
executionStrategyCtx.onDispatched();

futures.await().whenComplete((completeValueInfos, throwable) -> {
executionContext.running();
List<String> fieldsExecutedOnInitialResult = deferredExecutionSupport.getNonDeferredFieldNames(fieldNames);

BiConsumer<List<Object>, Throwable> handleResultsConsumer = handleResults(executionContext, fieldsExecutedOnInitialResult, overallResult);
if (throwable != null) {
handleResultsConsumer.accept(null, throwable.getCause());
executionContext.finished();
return;
}

Expand All @@ -75,17 +79,21 @@ public CompletableFuture<ExecutionResult> execute(ExecutionContext executionCont
dataLoaderDispatcherStrategy.executionStrategyOnFieldValuesInfo(completeValueInfos);
executionStrategyCtx.onFieldValuesInfo(completeValueInfos);
fieldValuesFutures.await().whenComplete(handleResultsConsumer);
executionContext.finished();
}).exceptionally((ex) -> {
executionContext.running();
// if there are any issues with combining/handling the field results,
// complete the future at all costs and bubble up any thrown exception so
// the execution does not hang.
dataLoaderDispatcherStrategy.executionStrategyOnFieldValuesException(ex);
executionStrategyCtx.onFieldValuesException();
overallResult.completeExceptionally(ex);
executionContext.finished();
return null;
});

overallResult.whenComplete(executionStrategyCtx::onCompleted);
executionContext.finished();
return overallResult;
}
}
4 changes: 3 additions & 1 deletion src/main/java/graphql/execution/DataLoaderDispatchStrategy.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,10 @@

import graphql.Internal;
import graphql.schema.DataFetcher;
import graphql.schema.DataFetchingEnvironment;

import java.util.List;
import java.util.function.Supplier;

@Internal
public interface DataLoaderDispatchStrategy {
Expand Down Expand Up @@ -44,7 +46,7 @@ default void executeObjectOnFieldValuesException(Throwable t, ExecutionStrategyP
default void fieldFetched(ExecutionContext executionContext,
ExecutionStrategyParameters executionStrategyParameters,
DataFetcher<?> dataFetcher,
Object fetchedValue) {
Object fetchedValue, Supplier<DataFetchingEnvironment> dataFetchingEnvironment) {

}

Expand Down
6 changes: 5 additions & 1 deletion src/main/java/graphql/execution/Execution.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,8 +29,8 @@
import graphql.schema.GraphQLObjectType;
import graphql.schema.GraphQLSchema;
import graphql.schema.impl.SchemaUtil;
import org.jspecify.annotations.NonNull;
import graphql.util.FpKit;
import org.jspecify.annotations.NonNull;
import org.reactivestreams.Publisher;

import java.util.Collections;
Expand Down Expand Up @@ -58,6 +58,8 @@ public class Execution {
private final ValueUnboxer valueUnboxer;
private final boolean doNotAutomaticallyDispatchDataLoader;

public static final String EXECUTION_CONTEXT_KEY = "__GraphQL_Java_ExecutionContext";

public Execution(ExecutionStrategy queryStrategy,
ExecutionStrategy mutationStrategy,
ExecutionStrategy subscriptionStrategy,
Expand Down Expand Up @@ -114,6 +116,8 @@ public CompletableFuture<ExecutionResult> execute(Document document, GraphQLSche
.build();

executionContext.getGraphQLContext().put(ResultNodesInfo.RESULT_NODES_INFO, executionContext.getResultNodesInfo());
executionContext.getGraphQLContext().put(EXECUTION_CONTEXT_KEY, executionContext);


InstrumentationExecutionParameters parameters = new InstrumentationExecutionParameters(
executionInput, graphQLSchema
Expand Down
15 changes: 15 additions & 0 deletions src/main/java/graphql/execution/ExecutionContext.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
import java.util.Locale;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Consumer;
import java.util.function.Supplier;
Expand Down Expand Up @@ -63,6 +64,8 @@ public class ExecutionContext {
private final Supplier<ExecutableNormalizedOperation> queryTree;
private final boolean propagateErrorsOnNonNullContractFailure;

private final AtomicInteger isRunning = new AtomicInteger(0);

// this is modified after creation so it needs to be volatile to ensure visibility across Threads
private volatile DataLoaderDispatchStrategy dataLoaderDispatcherStrategy = DataLoaderDispatchStrategy.NO_OP;

Expand Down Expand Up @@ -349,4 +352,16 @@ public ExecutionContext transform(Consumer<ExecutionContextBuilder> builderConsu
public ResultNodesInfo getResultNodesInfo() {
return resultNodesInfo;
}

public boolean isRunning() {
return isRunning.get() > 0;
}

public void running() {
isRunning.incrementAndGet();
}

public void finished() {
isRunning.decrementAndGet();
}
}
27 changes: 24 additions & 3 deletions src/main/java/graphql/execution/ExecutionStrategy.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -201,6 +201,7 @@ public static String mkNameForPath(List<Field> currentField) {
@SuppressWarnings("unchecked")
@DuckTyped(shape = "CompletableFuture<Map<String, Object>> | Map<String, Object>")
protected Object executeObject(ExecutionContext executionContext, ExecutionStrategyParameters parameters) throws NonNullableFieldWasNullException {
executionContext.running();
DataLoaderDispatchStrategy dataLoaderDispatcherStrategy = executionContext.getDataLoaderDispatcherStrategy();
dataLoaderDispatcherStrategy.executeObject(executionContext, parameters);
Instrumentation instrumentation = executionContext.getInstrumentation();
Expand All @@ -225,15 +226,18 @@ protected Object executeObject(ExecutionContext executionContext, ExecutionStrat
if (fieldValueInfosResult instanceof CompletableFuture) {
CompletableFuture<List<FieldValueInfo>> fieldValueInfos = (CompletableFuture<List<FieldValueInfo>>) fieldValueInfosResult;
fieldValueInfos.whenComplete((completeValueInfos, throwable) -> {
executionContext.running();
if (throwable != null) {
handleResultsConsumer.accept(null, throwable);
executionContext.finished();
return;
}

Async.CombinedBuilder<Object> resultFutures = fieldValuesCombinedBuilder(completeValueInfos);
dataLoaderDispatcherStrategy.executeObjectOnFieldValuesInfo(completeValueInfos, parameters);
resolveObjectCtx.onFieldValuesInfo(completeValueInfos);
resultFutures.await().whenComplete(handleResultsConsumer);
executionContext.finished();
}).exceptionally((ex) -> {
// if there are any issues with combining/handling the field results,
// complete the future at all costs and bubble up any thrown exception so
Expand All @@ -244,6 +248,7 @@ protected Object executeObject(ExecutionContext executionContext, ExecutionStrat
return null;
});
overallResult.whenComplete(resolveObjectCtx::onCompleted);
executionContext.finished();
return overallResult;
} else {
List<FieldValueInfo> completeValueInfos = (List<FieldValueInfo>) fieldValueInfosResult;
Expand All @@ -257,10 +262,12 @@ protected Object executeObject(ExecutionContext executionContext, ExecutionStrat
CompletableFuture<List<Object>> completedValues = (CompletableFuture<List<Object>>) completedValuesObject;
completedValues.whenComplete(handleResultsConsumer);
overallResult.whenComplete(resolveObjectCtx::onCompleted);
executionContext.finished();
return overallResult;
} else {
Map<String, Object> fieldValueMap = buildFieldValueMap(fieldsExecutedOnInitialResult, (List<Object>) completedValuesObject);
resolveObjectCtx.onCompleted(fieldValueMap, null);
executionContext.finished();
return fieldValueMap;
}
}
Expand All @@ -276,12 +283,15 @@ protected Object executeObject(ExecutionContext executionContext, ExecutionStrat

private BiConsumer<List<Object>, Throwable> buildFieldValueMap(List<String> fieldNames, CompletableFuture<Map<String, Object>> overallResult, ExecutionContext executionContext) {
return (List<Object> results, Throwable exception) -> {
executionContext.running();
if (exception != null) {
handleValueException(overallResult, exception, executionContext);
executionContext.finished();
return;
}
Map<String, Object> resolvedValuesByField = buildFieldValueMap(fieldNames, results);
overallResult.complete(resolvedValuesByField);
executionContext.finished();
};
}

Expand Down Expand Up @@ -496,7 +506,7 @@ private Object fetchField(GraphQLFieldDefinition fieldDef, ExecutionContext exec
dataFetcher = instrumentation.instrumentDataFetcher(dataFetcher, instrumentationFieldFetchParams, executionContext.getInstrumentationState());
dataFetcher = executionContext.getDataLoaderDispatcherStrategy().modifyDataFetcher(dataFetcher);
Object fetchedObject = invokeDataFetcher(executionContext, parameters, fieldDef, dataFetchingEnvironment, dataFetcher);
executionContext.getDataLoaderDispatcherStrategy().fieldFetched(executionContext, parameters, dataFetcher, fetchedObject);
executionContext.getDataLoaderDispatcherStrategy().fieldFetched(executionContext, parameters, dataFetcher, fetchedObject, dataFetchingEnvironment);
fetchCtx.onDispatched();
fetchCtx.onFetchedValue(fetchedObject);
// if it's a subscription, leave any reactive objects alone
Expand All @@ -509,11 +519,15 @@ private Object fetchField(GraphQLFieldDefinition fieldDef, ExecutionContext exec
CompletableFuture<Object> fetchedValue = (CompletableFuture<Object>) fetchedObject;
return fetchedValue
.handle((result, exception) -> {
executionContext.running();
fetchCtx.onCompleted(result, exception);
if (exception != null) {
return handleFetchingException(dataFetchingEnvironment.get(), parameters, exception);
CompletableFuture<Object> handleFetchingExceptionResult = handleFetchingException(dataFetchingEnvironment.get(), parameters, exception);
executionContext.finished();
return handleFetchingExceptionResult;
} else {
// we can simply return the fetched value CF and avoid a allocation
executionContext.finished();
return fetchedValue;
}
})
Expand Down Expand Up @@ -553,7 +567,7 @@ protected Supplier<ExecutableNormalizedField> getNormalizedField(ExecutionContex
protected FetchedValue unboxPossibleDataFetcherResult(ExecutionContext executionContext,
ExecutionStrategyParameters parameters,
Object result) {

executionContext.running();
if (result instanceof DataFetcherResult) {
DataFetcherResult<?> dataFetcherResult = (DataFetcherResult<?>) result;

Expand All @@ -567,9 +581,11 @@ protected FetchedValue unboxPossibleDataFetcherResult(ExecutionContext execution
localContext = parameters.getLocalContext();
}
Object unBoxedValue = executionContext.getValueUnboxer().unbox(dataFetcherResult.getData());
executionContext.finished();
return new FetchedValue(unBoxedValue, dataFetcherResult.getErrors(), localContext);
} else {
Object unBoxedValue = executionContext.getValueUnboxer().unbox(result);
executionContext.finished();
return new FetchedValue(unBoxedValue, ImmutableList.of(), parameters.getLocalContext());
}
}
Expand Down Expand Up @@ -638,6 +654,7 @@ protected FieldValueInfo completeField(ExecutionContext executionContext, Execut
}

private FieldValueInfo completeField(GraphQLFieldDefinition fieldDef, ExecutionContext executionContext, ExecutionStrategyParameters parameters, FetchedValue fetchedValue) {
executionContext.running();
GraphQLObjectType parentType = (GraphQLObjectType) parameters.getExecutionStepInfo().getUnwrappedNonNullType();
ExecutionStepInfo executionStepInfo = createExecutionStepInfo(executionContext, parameters, fieldDef, parentType);

Expand All @@ -661,6 +678,7 @@ private FieldValueInfo completeField(GraphQLFieldDefinition fieldDef, ExecutionC
CompletableFuture<Object> executionResultFuture = fieldValueInfo.getFieldValueFuture();
ctxCompleteField.onDispatched();
executionResultFuture.whenComplete(ctxCompleteField::onCompleted);
executionContext.finished();
return fieldValueInfo;
}

Expand Down Expand Up @@ -833,13 +851,16 @@ protected FieldValueInfo completeValueForList(ExecutionContext executionContext,
overallResult.whenComplete(completeListCtx::onCompleted);

resultsFuture.whenComplete((results, exception) -> {
executionContext.running();
if (exception != null) {
executionContext.finished();
handleValueException(overallResult, exception, executionContext);
return;
}
List<Object> completedResults = new ArrayList<>(results.size());
completedResults.addAll(results);
overallResult.complete(completedResults);
executionContext.finished();
});
listOrPromiseToList = overallResult;
} else {
Expand Down
Loading
Loading