Support Callable as a controller method return value

Closes gh-316

Author: rstoyanchev

spring-graphql-docs/src/docs/asciidoc/index.adoc

@@ -1227,8 +1227,13 @@ See <<controllers-schema-mapping-data-loader>>.
 
 |===
 
-Schema mapping handler methods can return any value, including Reactor `Mono` and
-`Flux` as described in <<execution-reactive-datafetcher>>.
+Schema mapping handler methods can return:
+
+- A resolved value of any type.
+- `Mono` and `Flux` for asynchronous value(s). Supported for controller methods and for
+  any `DataFetcher` as described in <<execution-reactive-datafetcher>>.
+- `java.util.concurrent.Callable` to have the value(s) produced asynchronously.
+  For this to work, `AnnotatedControllerConfigurer` must be configured with an `Executor`.
 
 
 
@@ -1574,6 +1579,10 @@ Batch mapping methods can return:
 | `Map<K,V>`, `Collection<V>`
 | Imperative variants, e.g. without remote calls to make.
 
+| `Callable<Map<K,V>>`, `Callable<Collection<V>>`
+| Imperative variants to be invoked asynchronously. For this to work,
+  `AnnotatedControllerConfigurer` must be configured with an `Executor`.
+
 |===
 
 

spring-graphql/src/main/java/org/springframework/graphql/data/method/InvocableHandlerMethodSupport.java

@@ -19,14 +19,20 @@
 import java.lang.reflect.Method;
 import java.util.Arrays;
 import java.util.List;
+import java.util.concurrent.Callable;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.Executor;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
+import graphql.GraphQLContext;
 import reactor.core.publisher.Mono;
 
 import org.springframework.core.CoroutinesUtils;
 import org.springframework.core.KotlinDetector;
+import org.springframework.graphql.execution.ReactorContextManager;
 import org.springframework.lang.Nullable;
+import org.springframework.util.Assert;
 
 /**
  * Extension of {@link HandlerMethod} that adds support for invoking the
@@ -40,8 +46,24 @@ public abstract class InvocableHandlerMethodSupport extends HandlerMethod {
 	private static final Object NO_VALUE = new Object();
 
 
-	protected InvocableHandlerMethodSupport(HandlerMethod handlerMethod) {
+	private final boolean hasCallableReturnValue;
+
+	@Nullable
+	private final Executor executor;
+
+
+	/**
+	 * Create an instance.
+	 * @param handlerMethod the controller method
+	 * @param executor an {@link Executor} to use for {@link Callable} return values
+	 */
+	protected InvocableHandlerMethodSupport(HandlerMethod handlerMethod, @Nullable Executor executor) {
 		super(handlerMethod.createWithResolvedBean());
+		this.hasCallableReturnValue = getReturnType().getParameterType().equals(Callable.class);
+		this.executor = executor;
+		Assert.isTrue(!this.hasCallableReturnValue || this.executor != null,
+				"Controller method declared with Callable return value, but no Executor configured: " +
+						handlerMethod.getBridgedMethod().toGenericString());
 	}
 
 
@@ -51,8 +73,9 @@ protected InvocableHandlerMethodSupport(HandlerMethod handlerMethod) {
 	 * @return the value returned from the method or a {@code Mono<Throwable>}
 	 * if the invocation fails.
 	 */
+	@SuppressWarnings("ReactiveStreamsUnusedPublisher")
 	@Nullable
-	protected Object doInvoke(Object... argValues) {
+	protected Object doInvoke(GraphQLContext graphQLContext, Object... argValues) {
 		if (logger.isTraceEnabled()) {
 			logger.trace("Arguments: " + Arrays.toString(argValues));
 		}
@@ -61,7 +84,8 @@ protected Object doInvoke(Object... argValues) {
 			if (KotlinDetector.isSuspendingFunction(method)) {
 				return CoroutinesUtils.invokeSuspendingFunction(method, getBean(), argValues);
 			}
-			return method.invoke(getBean(), argValues);
+			Object result = method.invoke(getBean(), argValues);
+			return handleReturnValue(graphQLContext, result);
 		}
 		catch (IllegalArgumentException ex) {
 			assertTargetBean(method, getBean(), argValues);
@@ -84,6 +108,24 @@ protected Object doInvoke(Object... argValues) {
 		}
 	}
 
+	@Nullable
+	private Object handleReturnValue(GraphQLContext graphQLContext, @Nullable Object result) {
+		if (this.hasCallableReturnValue && result != null) {
+			return CompletableFuture.supplyAsync(
+					() -> {
+						try {
+							return ReactorContextManager.invokeCallable((Callable<?>) result, graphQLContext);
+						}
+						catch (Exception ex) {
+							throw new IllegalStateException(
+									"Failure in Callable returned from " + getBridgedMethod().toGenericString(), ex);
+						}
+					},
+					this.executor);
+		}
+		return result;
+	}
+
 	/**
 	 * Use this method to resolve the arguments asynchronously. This is only
 	 * useful when at least one of the values is a {@link Mono}

spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/AnnotatedControllerConfigurer.java

@@ -25,6 +25,8 @@
 import java.util.LinkedHashSet;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.Callable;
+import java.util.concurrent.Executor;
 import java.util.stream.Collectors;
 
 import javax.validation.Validator;
@@ -110,6 +112,12 @@ public class AnnotatedControllerConfigurer
 			"javax.validation.executable.ExecutableValidator",
 			AnnotatedControllerConfigurer.class.getClassLoader());
 
+
+	private final FormattingConversionService conversionService = new DefaultFormattingConversionService();
+
+	@Nullable
+	private Executor executor;
+
 	@Nullable
 	private ApplicationContext applicationContext;
 
@@ -119,8 +127,6 @@ public class AnnotatedControllerConfigurer
 	@Nullable
 	private HandlerMethodInputValidator validator;
 
-	private FormattingConversionService conversionService = new DefaultFormattingConversionService();
-
 
 	/**
 	 * Add a {@code FormatterRegistrar} to customize the {@link ConversionService}
@@ -132,6 +138,17 @@ public void addFormatterRegistrar(FormatterRegistrar registrar) {
 		registrar.registerFormatters(this.conversionService);
 	}
 
+	/**
+	 * Configure an {@link Executor} to use for asynchronous handling of
+	 * {@link Callable} return values from controller methods.
+	 * <p>By default, this is not set in which case controller methods with a
+	 * {@code Callable} return value cannot be registered.
+	 * @param executor the executor to use
+	 */
+	public void setExecutor(Executor executor) {
+		this.executor = executor;
+	}
+
 	@Override
 	public void setApplicationContext(ApplicationContext applicationContext) {
 		this.applicationContext = applicationContext;
@@ -195,7 +212,7 @@ public void configure(RuntimeWiring.Builder runtimeWiringBuilder) {
 		findHandlerMethods().forEach((info) -> {
 			DataFetcher<?> dataFetcher;
 			if (!info.isBatchMapping()) {
-				dataFetcher = new SchemaMappingDataFetcher(info, this.argumentResolvers, this.validator);
+				dataFetcher = new SchemaMappingDataFetcher(info, this.argumentResolvers, this.validator, this.executor);
 			}
 			else {
 				String dataLoaderKey = registerBatchLoader(info);
@@ -359,13 +376,16 @@ private String registerBatchLoader(MappingInfo info) {
 		BatchLoaderRegistry registry = obtainApplicationContext().getBean(BatchLoaderRegistry.class);
 
 		HandlerMethod handlerMethod = info.getHandlerMethod();
-		BatchLoaderHandlerMethod invocable = new BatchLoaderHandlerMethod(handlerMethod);
+		BatchLoaderHandlerMethod invocable = new BatchLoaderHandlerMethod(handlerMethod, this.executor);
 
-		Class<?> clazz = handlerMethod.getReturnType().getParameterType();
-		if (clazz.equals(Flux.class) || Collection.class.isAssignableFrom(clazz)) {
+		MethodParameter returnType = handlerMethod.getReturnType();
+		Class<?> clazz = returnType.getParameterType();
+		Class<?> nestedClass = (clazz.equals(Callable.class) ? returnType.nested().getNestedParameterType() : clazz);
+
+		if (clazz.equals(Flux.class) || Collection.class.isAssignableFrom(nestedClass)) {
 			registry.forName(dataLoaderKey).registerBatchLoader(invocable::invokeForIterable);
 		}
-		else if (clazz.equals(Mono.class) || clazz.equals(Map.class)) {
+		else if (clazz.equals(Mono.class) || nestedClass.equals(Map.class)) {
 			registry.forName(dataLoaderKey).registerMappedBatchLoader(invocable::invokeForMap);
 		}
 		else {
@@ -425,7 +445,7 @@ public HandlerMethod getHandlerMethod() {
 
 		@Override
 		public String toString() {
-			return this.coordinates + " -> " + this.handlerMethod.toString();
+			return this.coordinates + " -> " + this.handlerMethod;
 		}
 	}
 
@@ -442,25 +462,23 @@ static class SchemaMappingDataFetcher implements DataFetcher<Object> {
 		@Nullable
 		private final HandlerMethodInputValidator validator;
 
+		@Nullable
+		private final Executor executor;
+
 		private final boolean subscription;
 
 		public SchemaMappingDataFetcher(
 				MappingInfo info, HandlerMethodArgumentResolverComposite resolvers,
-				@Nullable HandlerMethodInputValidator validator) {
+				@Nullable HandlerMethodInputValidator validator,
+				@Nullable Executor executor) {
 
 			this.info = info;
 			this.argumentResolvers = resolvers;
 			this.validator = validator;
+			this.executor = executor;
 			this.subscription = this.info.getCoordinates().getTypeName().equalsIgnoreCase("Subscription");
 		}
 
-		/**
-		 * Return the {@link FieldCoordinates} the HandlerMethod is mapped to.
-		 */
-		public FieldCoordinates getCoordinates() {
-			return this.info.getCoordinates();
-		}
-
 		/**
 		 * Return the {@link HandlerMethod} used to fetch data.
 		 */
@@ -474,7 +492,7 @@ public HandlerMethod getHandlerMethod() {
 		public Object get(DataFetchingEnvironment environment) throws Exception {
 
 			DataFetcherHandlerMethod handlerMethod = new DataFetcherHandlerMethod(
-					getHandlerMethod(), this.argumentResolvers, this.validator, this.subscription);
+					getHandlerMethod(), this.argumentResolvers, this.validator, this.executor, this.subscription);
 
 			return handlerMethod.invoke(environment);
 		}

spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/BatchLoaderHandlerMethod.java

@@ -19,6 +19,9 @@
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Map;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.Executor;
+import java.util.function.Function;
 
 import graphql.GraphQLContext;
 import org.dataloader.BatchLoaderEnvironment;
@@ -50,8 +53,8 @@ public class BatchLoaderHandlerMethod extends InvocableHandlerMethodSupport {
 			AnnotatedControllerConfigurer.class.getClassLoader());
 
 
-	public BatchLoaderHandlerMethod(HandlerMethod handlerMethod) {
-		super(handlerMethod);
+	public BatchLoaderHandlerMethod(HandlerMethod handlerMethod, @Nullable Executor executor) {
+		super(handlerMethod, executor);
 	}
 
 
@@ -69,11 +72,11 @@ public BatchLoaderHandlerMethod(HandlerMethod handlerMethod) {
 	public <K, V> Mono<Map<K, V>> invokeForMap(Collection<K> keys, BatchLoaderEnvironment environment) {
 		Object[] args = getMethodArgumentValues(keys, environment);
 		if (doesNotHaveAsyncArgs(args)) {
-			Object result = doInvoke(args);
+			Object result = doInvoke(environment.getContext(), args);
 			return toMonoMap(result);
 		}
 		return toArgsMono(args).flatMap(argValues -> {
-			Object result = doInvoke(argValues);
+			Object result = doInvoke(environment.getContext(), argValues);
 			return toMonoMap(result);
 		});
 	}
@@ -90,11 +93,11 @@ public <K, V> Mono<Map<K, V>> invokeForMap(Collection<K> keys, BatchLoaderEnviro
 	public <V> Flux<V> invokeForIterable(Collection<?> keys, BatchLoaderEnvironment environment) {
 		Object[] args = getMethodArgumentValues(keys, environment);
 		if (doesNotHaveAsyncArgs(args)) {
-			Object result = doInvoke(args);
+			Object result = doInvoke(environment.getContext(), args);
 			return toFlux(result);
 		}
 		return toArgsMono(args).flatMapMany(resolvedArgs -> {
-			Object result = doInvoke(resolvedArgs);
+			Object result = doInvoke(environment.getContext(), resolvedArgs);
 			return toFlux(result);
 		});
 	}
@@ -165,6 +168,9 @@ private static <K, V> Mono<Map<K, V>> toMonoMap(@Nullable Object result) {
 		else if (result instanceof Mono) {
 			return (Mono<Map<K, V>>) result;
 		}
+		else if (result instanceof CompletableFuture) {
+			return Mono.fromFuture((CompletableFuture<? extends Map<K,V>>) result);
+		}
 		return Mono.error(new IllegalStateException("Unexpected return value: " + result));
 	}
 
@@ -176,6 +182,10 @@ private static <V> Flux<V> toFlux(@Nullable Object result) {
 		else if (result instanceof Flux) {
 			return (Flux<V>) result;
 		}
+		else if (result instanceof CompletableFuture) {
+			return Mono.fromFuture((CompletableFuture<? extends Collection<V>>) result)
+					.flatMapIterable(Function.identity());
+		}
 		return Flux.error(new IllegalStateException("Unexpected return value: " + result));
 	}
 

spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/DataFetcherHandlerMethod.java

@@ -16,6 +16,7 @@
 package org.springframework.graphql.data.method.annotation.support;
 
 import java.util.Arrays;
+import java.util.concurrent.Executor;
 
 import graphql.schema.DataFetchingEnvironment;
 import org.reactivestreams.Publisher;
@@ -65,9 +66,9 @@ public class DataFetcherHandlerMethod extends InvocableHandlerMethodSupport {
 	 */
 	public DataFetcherHandlerMethod(HandlerMethod handlerMethod,
 			HandlerMethodArgumentResolverComposite resolvers, @Nullable HandlerMethodInputValidator validator,
-			boolean subscription) {
+			@Nullable Executor executor, boolean subscription) {
 
-		super(handlerMethod);
+		super(handlerMethod, executor);
 		Assert.isTrue(!resolvers.getResolvers().isEmpty(), "No argument resolvers");
 		this.resolvers = resolvers;
 		this.validator = validator;
@@ -118,17 +119,17 @@ public Object invoke(DataFetchingEnvironment environment) {
 		}
 
 		if (Arrays.stream(args).noneMatch(arg -> arg instanceof Mono)) {
-			return validateAndInvoke(args);
+			return validateAndInvoke(args, environment);
 		}
 
 		return this.subscription ?
 				toArgsMono(args).flatMapMany(argValues -> {
-					Object result = validateAndInvoke(argValues);
+					Object result = validateAndInvoke(argValues, environment);
 					Assert.state(result instanceof Publisher, "Expected a Publisher from a Subscription response");
 					return Flux.from((Publisher<?>) result);
 				}) :
 				toArgsMono(args).flatMap(argValues -> {
-					Object result = validateAndInvoke(argValues);
+					Object result = validateAndInvoke(argValues, environment);
 					if (result instanceof Mono) {
 						return (Mono<?>) result;
 					}
@@ -183,11 +184,11 @@ private Object[] getMethodArgumentValues(
 	}
 
 	@Nullable
-	private Object validateAndInvoke(Object[] args) {
+	private Object validateAndInvoke(Object[] args, DataFetchingEnvironment environment) {
 		if (this.validator != null) {
 			this.validator.validate(this, args);
 		}
-		return doInvoke(args);
+		return doInvoke(environment.getGraphQlContext(), args);
 	}
 
 }

spring-graphql/src/main/java/org/springframework/graphql/execution/ContextDataFetcherDecorator.java

@@ -59,16 +59,11 @@ private ContextDataFetcherDecorator(DataFetcher<?> delegate, boolean subscriptio
 
 	@Override
 	public Object get(DataFetchingEnvironment environment) throws Exception {
-		ContextView contextView = ReactorContextManager.getReactorContext(environment.getGraphQlContext());
 
-		Object value;
-		try {
-			ReactorContextManager.restoreThreadLocalValues(contextView);
-			value = this.delegate.get(environment);
-		}
-		finally {
-			ReactorContextManager.resetThreadLocalValues(contextView);
-		}
+		Object value = ReactorContextManager.invokeCallable(() ->
+				this.delegate.get(environment), environment.getGraphQlContext());
+
+		ContextView contextView = ReactorContextManager.getReactorContext(environment.getGraphQlContext());
 
 		if (this.subscription) {
 			return (!contextView.isEmpty() ? Flux.from((Publisher<?>) value).contextWrite(contextView) : value);