graphql.execution.instrumentation.threadpools.ExecutorInstrumentation Maven / Gradle / Ivy
package graphql.execution.instrumentation.threadpools;
import com.google.common.annotations.Beta;
import graphql.Assert;
import graphql.Internal;
import graphql.TrivialDataFetcher;
import graphql.execution.Async;
import graphql.execution.instrumentation.InstrumentationState;
import graphql.execution.instrumentation.SimplePerformantInstrumentation;
import graphql.execution.instrumentation.parameters.InstrumentationFieldFetchParameters;
import graphql.schema.DataFetcher;
import graphql.schema.DataFetchingEnvironment;
import org.jetbrains.annotations.NotNull;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.Executor;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.function.Supplier;
import static graphql.execution.instrumentation.threadpools.ExecutorInstrumentation.Action.FETCHING;
import static graphql.execution.instrumentation.threadpools.ExecutorInstrumentation.Action.PROCESSING;
/**
* This instrumentation can be used to control on what thread calls to {@link DataFetcher}s happen on.
*
* If your data fetching is inherently IO bound then you could use a IO oriented thread pool for your fetches and transfer control
* back to a CPU oriented thread pool and allow graphql-java code to run the post-processing of results there.
*
* An IO oriented thread pool is typically a multiple of {@link Runtime#availableProcessors()} while a CPU oriented thread pool
* is typically no more than {@link Runtime#availableProcessors()}.
*
* The instrumentation will use the {@link graphql.execution.instrumentation.Instrumentation#instrumentDataFetcher(DataFetcher, InstrumentationFieldFetchParameters, InstrumentationState)}
* method to change your data fetchers, so they are executed on a thread pool dedicated to fetching (if you provide one).
*
* Once the data fetcher value is returns it will transfer control back to a processing thread pool (if you provide one).
*
* This code uses {@link CompletableFuture#supplyAsync(Supplier, Executor)} and {@link CompletableFuture#thenApplyAsync(Function, Executor)} to transfer
* control between thread pools.
*/
@Internal
@Beta
public class ExecutorInstrumentation extends SimplePerformantInstrumentation {
private static final Consumer NOOP = a -> {
};
/**
* This describes what action is currently being done. This is mostly intended for testing.
*/
enum Action {FETCHING, PROCESSING}
private final Executor fetchExecutor;
private final Executor processingExecutor;
private final Consumer actionObserver;
private ExecutorInstrumentation(Executor fetchExecutor, Executor processingExecutor, Consumer actionObserver) {
this.fetchExecutor = fetchExecutor;
this.processingExecutor = processingExecutor;
this.actionObserver = actionObserver;
}
public Executor getFetchExecutor() {
return fetchExecutor;
}
public Executor getProcessingExecutor() {
return processingExecutor;
}
public static Builder newThreadPoolExecutionInstrumentation() {
return new Builder();
}
public static class Builder {
Executor fetchExecutor;
Executor processingExecutor;
private Consumer actionObserver;
public Builder fetchExecutor(Executor fetchExecutor) {
this.fetchExecutor = fetchExecutor;
return this;
}
public Builder processingExecutor(Executor processingExecutor) {
this.processingExecutor = processingExecutor;
return this;
}
/**
* This is really intended for testing but this consumer will be called during
* stages to indicate what is happening.
*
* @param actionObserver the observer code
*
* @return this builder
*/
public Builder actionObserver(Consumer actionObserver) {
this.actionObserver = Assert.assertNotNull(actionObserver);
return this;
}
public ExecutorInstrumentation build() {
return new ExecutorInstrumentation(fetchExecutor, processingExecutor, actionObserver != null ? actionObserver : NOOP);
}
}
@Override
public @NotNull DataFetcher instrumentDataFetcher(DataFetcher originalDataFetcher, InstrumentationFieldFetchParameters parameters, InstrumentationState state) {
if (originalDataFetcher instanceof TrivialDataFetcher) {
return originalDataFetcher;
}
return environment -> {
CompletableFuture> invokedCF;
if (fetchExecutor != null) {
// run the fetch asynchronously via the fetch executor
// the CF will be left running on that fetch executors thread
invokedCF = CompletableFuture.supplyAsync(invokedAsync(originalDataFetcher, environment), fetchExecutor);
} else {
invokedCF = invokedSync(originalDataFetcher, environment);
}
if (processingExecutor != null) {
invokedCF = invokedCF.thenApplyAsync(processingControl(), processingExecutor);
} else {
invokedCF = invokedCF.thenApply(processingControl());
}
return invokedCF.thenCompose(cs -> cs);
};
}
private Supplier> invokedAsync(DataFetcher originalDataFetcher, DataFetchingEnvironment environment) {
return () -> {
actionObserver.accept(FETCHING);
return invokeOriginalDF(originalDataFetcher, environment);
};
}
private CompletableFuture> invokedSync(DataFetcher originalDataFetcher, DataFetchingEnvironment environment) {
actionObserver.accept(FETCHING);
return CompletableFuture.completedFuture(invokeOriginalDF(originalDataFetcher, environment));
}
private Function, CompletionStage> processingControl() {
return completionStage -> {
actionObserver.accept(PROCESSING);
return completionStage;
};
}
private CompletionStage invokeOriginalDF(DataFetcher originalDataFetcher, DataFetchingEnvironment environment) {
Object value;
try {
value = originalDataFetcher.get(environment);
} catch (Exception e) {
return Async.exceptionallyCompletedFuture(e);
}
if (value instanceof CompletionStage) {
return ((CompletionStage) value);
} else {
return CompletableFuture.completedFuture(value);
}
}
}