Support Callable as a controller method return value

Closes gh-316

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.InvocationTargetException;
 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 @@ public abstract class InvocableHandlerMethodSupport extends 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 @@ public abstract class InvocableHandlerMethodSupport extends HandlerMethod {
 			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 @@ public abstract class InvocableHandlerMethodSupport extends HandlerMethod {
 		}
 	}
 
+	@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.HashMap;
 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 class AnnotatedControllerConfigurer
 		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 class AnnotatedControllerConfigurer
 		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 @@ public class AnnotatedControllerConfigurer
 		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 class AnnotatedControllerConfigurer
 
 		@Override
 		public String toString() {
-			return this.coordinates + " -> " + this.handlerMethod.toString();
+			return this.coordinates + " -> " + this.handlerMethod;
 		}
 	}
 
@@ -442,25 +462,23 @@ public class AnnotatedControllerConfigurer
 		@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 class AnnotatedControllerConfigurer
 		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.security.Principal;
 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 class BatchLoaderHandlerMethod extends InvocableHandlerMethodSupport {
 	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 class BatchLoaderHandlerMethod extends InvocableHandlerMethodSupport {
 	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 @@ public class BatchLoaderHandlerMethod extends InvocableHandlerMethodSupport {
 		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 @@ public class BatchLoaderHandlerMethod extends InvocableHandlerMethodSupport {
 		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 class DataFetcherHandlerMethod extends InvocableHandlerMethodSupport {
 		}
 
 		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 @@ public class DataFetcherHandlerMethod extends InvocableHandlerMethodSupport {
 	}
 
 	@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 @@ final class ContextDataFetcherDecorator implements DataFetcher<Object> {
 
 	@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);

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

@@ -18,6 +18,7 @@ package org.springframework.graphql.execution;
 
 import java.util.LinkedHashMap;
 import java.util.Map;
+import java.util.concurrent.Callable;
 
 import graphql.GraphQLContext;
 import reactor.util.context.Context;
@@ -70,7 +71,7 @@ public abstract class ReactorContextManager {
 	 * Use the given accessor to extract ThreadLocal values and save them in a
 	 * sub-map in the given {@link Context}, so those can be restored later
 	 * around the execution of data fetchers and exception resolvers. The accessor
-	 * instance is also saved in the Reactor Context so it can be used to
+	 * instance is also saved in the Reactor Context, so it can be used to
 	 * actually restore and reset ThreadLocal values.
 	 * @param accessor the accessor to use
 	 * @param context the context to write to if there are ThreadLocal values
@@ -89,6 +90,24 @@ public abstract class ReactorContextManager {
 				THREAD_ID, Thread.currentThread().getId()));
 	}
 
+	/**
+	 * Restore {@code ThreadLocal} values, invoke the given {@code Callable},
+	 * and reset the {@code ThreadLocal} values.
+	 * @param callable the callable to invoke
+	 * @param graphQlContext the current {@code GraphQLContext}
+	 * @return the return value from the invocation
+	 */
+	public static <T> T invokeCallable(Callable<T> callable, GraphQLContext graphQlContext) throws Exception {
+		ContextView contextView = getReactorContext(graphQlContext);
+		try {
+			ReactorContextManager.restoreThreadLocalValues(contextView);
+			return callable.call();
+		}
+		finally {
+			ReactorContextManager.resetThreadLocalValues(contextView);
+		}
+	}
+
 	/**
 	 * Look up saved ThreadLocal values and restore them if any are found.
 	 * This is a no-op if invoked on the thread that values were extracted on.

spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/BatchMappingDetectionTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2021 the original author or authors.
+ * Copyright 2002-2022 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -17,6 +17,7 @@ package org.springframework.graphql.data.method.annotation.support;
 
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.Callable;
 
 import graphql.GraphQLContext;
 import graphql.schema.DataFetcher;
@@ -28,6 +29,7 @@ import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
 import org.springframework.context.annotation.AnnotationConfigApplicationContext;
+import org.springframework.core.task.SimpleAsyncTaskExecutor;
 import org.springframework.graphql.Author;
 import org.springframework.graphql.Book;
 import org.springframework.graphql.data.method.annotation.BatchMapping;
@@ -59,13 +61,14 @@ public class BatchMappingDetectionTests {
 
 		assertThat(dataFetcherMap).containsOnlyKeys("Book");
 		assertThat(dataFetcherMap.get("Book")).containsOnlyKeys(
-				"authorFlux", "authorList", "authorMonoMap", "authorMap", "authorEnvironment");
+				"authorFlux", "authorList", "authorMonoMap", "authorMap", "authorCallableMap", "authorEnvironment");
 
 		DataLoaderRegistry registry = new DataLoaderRegistry();
 		this.batchLoaderRegistry.registerDataLoaders(registry, GraphQLContext.newContext().build());
 
 		assertThat(registry.getDataLoadersMap()).containsOnlyKeys(
-				"Book.authorFlux", "Book.authorList", "Book.authorMonoMap", "Book.authorMap", "Book.authorEnvironment");
+				"Book.authorFlux", "Book.authorList", "Book.authorMonoMap", "Book.authorMap",
+				"Book.authorCallableMap", "Book.authorEnvironment");
 	}
 
 	@Test
@@ -87,6 +90,7 @@ public class BatchMappingDetectionTests {
 		context.refresh();
 
 		AnnotatedControllerConfigurer configurer = new AnnotatedControllerConfigurer();
+		configurer.setExecutor(new SimpleAsyncTaskExecutor());
 		configurer.setApplicationContext(context);
 		configurer.afterPropertiesSet();
 
@@ -120,6 +124,11 @@ public class BatchMappingDetectionTests {
 			return null;
 		}
 
+		@BatchMapping
+		public Callable<Map<Book, Author>> authorCallableMap(List<Book> books) {
+			return null;
+		}
+
 		@BatchMapping
 		public List<Author> authorEnvironment(BatchLoaderEnvironment environment, List<Book> books) {
 			return null;

spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/BatchMappingInvocationTests.java

@@ -18,6 +18,7 @@ package org.springframework.graphql.data.method.annotation.support;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.Callable;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -51,7 +52,8 @@ public class BatchMappingInvocationTests extends BatchMappingTestSupport {
 				arguments(named("Returning Mono<Map<K,V>>", new BatchMonoMapController())),
 				arguments(named("Returning Map<K,V>", new BatchMapController())),
 				arguments(named("Returning Flux<V>", new BatchFluxController())),
-				arguments(named("Returning List<V>", new BatchListController()))
+				arguments(named("Returning List<V>", new BatchListController())),
+				arguments(named("Returning Callable<Map<K,V>>", new BatchCallableMapController()))
 		);
 	}
 
@@ -141,6 +143,7 @@ public class BatchMappingInvocationTests extends BatchMappingTestSupport {
 		}
 	}
 
+
 	@Controller
 	private static class BatchMapController extends CourseController {
 
@@ -153,8 +156,10 @@ public class BatchMappingInvocationTests extends BatchMappingTestSupport {
 		public Map<Course, List<Person>> students(List<Course> courses) {
 			return courses.stream().collect(Collectors.toMap(Function.identity(), Course::students));
 		}
+
 	}
 
+
 	@Controller
 	private static class BatchFluxController extends CourseController {
 
@@ -169,6 +174,7 @@ public class BatchMappingInvocationTests extends BatchMappingTestSupport {
 		}
 	}
 
+
 	@Controller
 	private static class BatchListController extends CourseController {
 
@@ -183,4 +189,21 @@ public class BatchMappingInvocationTests extends BatchMappingTestSupport {
 		}
 	}
 
+
+	@Controller
+	private static class BatchCallableMapController extends CourseController {
+
+		@BatchMapping
+		public Callable<Map<Course, Person>> instructor(List<Course> courses) {
+			return () -> courses.stream().collect(Collectors.toMap(Function.identity(), Course::instructor));
+		}
+
+		@BatchMapping
+		public Callable<Map<Course, List<Person>>> students(List<Course> courses) {
+			return () -> courses.stream().collect(Collectors.toMap(Function.identity(), Course::students));
+		}
+
+	}
+
+
 }

spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/BatchMappingTestSupport.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2021 the original author or authors.
+ * Copyright 2002-2022 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -28,6 +28,7 @@ import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
 
 import org.springframework.context.annotation.AnnotationConfigApplicationContext;
+import org.springframework.core.task.SimpleAsyncTaskExecutor;
 import org.springframework.graphql.ExecutionGraphQlService;
 import org.springframework.graphql.GraphQlSetup;
 import org.springframework.graphql.data.method.annotation.QueryMapping;
@@ -87,8 +88,13 @@ public class BatchMappingTestSupport {
 		context.registerBean(BatchLoaderRegistry.class, () -> registry);
 		context.refresh();
 
+		AnnotatedControllerConfigurer configurer = new AnnotatedControllerConfigurer();
+		configurer.setExecutor(new SimpleAsyncTaskExecutor());
+		configurer.setApplicationContext(context);
+		configurer.afterPropertiesSet();
+
 		return GraphQlSetup.schemaContent(schema)
-				.runtimeWiringForAnnotatedControllers(context)
+				.runtimeWiring(configurer)
 				.dataLoaders(registry)
 				.toGraphQlService();
 	}

spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/ContextValueMethodArgumentResolverTests.java

@@ -119,7 +119,7 @@ public class ContextValueMethodArgumentResolverTests {
 
 		DataFetcherHandlerMethod handlerMethod = new DataFetcherHandlerMethod(
 				new HandlerMethod(new TestController(), TestController.class.getMethod("handleMono", Mono.class)),
-				resolvers, null, false);
+				resolvers, null, null, false);
 
 		GraphQLContext graphQLContext = new GraphQLContext.Builder().build();
 

spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/DataFetcherHandlerMethodTests.java

@@ -0,0 +1,77 @@
+/*
+ * Copyright 2002-2022 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.graphql.data.method.annotation.support;
+
+
+import java.util.concurrent.Callable;
+import java.util.concurrent.CompletableFuture;
+
+import graphql.GraphQLContext;
+import graphql.schema.DataFetchingEnvironment;
+import graphql.schema.DataFetchingEnvironmentImpl;
+import org.junit.jupiter.api.Test;
+import org.mockito.Mockito;
+
+import org.springframework.core.task.SimpleAsyncTaskExecutor;
+import org.springframework.graphql.data.method.HandlerMethod;
+import org.springframework.graphql.data.method.HandlerMethodArgumentResolver;
+import org.springframework.graphql.data.method.HandlerMethodArgumentResolverComposite;
+import org.springframework.lang.Nullable;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * Unit tests for {@link DataFetcherHandlerMethod}.
+ *
+ * @author Rossen Stoyanchev
+ */
+public class DataFetcherHandlerMethodTests {
+
+
+	@Test
+	void callableReturnValue() throws Exception {
+
+		HandlerMethodArgumentResolverComposite resolvers = new HandlerMethodArgumentResolverComposite();
+		resolvers.addResolver(Mockito.mock(HandlerMethodArgumentResolver.class));
+
+		DataFetcherHandlerMethod handlerMethod = new DataFetcherHandlerMethod(
+				new HandlerMethod(new TestController(), TestController.class.getMethod("handleAndReturnCallable")),
+				resolvers, null, new SimpleAsyncTaskExecutor(), false);
+
+		GraphQLContext graphQLContext = new GraphQLContext.Builder().build();
+
+		DataFetchingEnvironment environment = DataFetchingEnvironmentImpl.newDataFetchingEnvironment()
+				.graphQLContext(graphQLContext)
+				.build();
+
+		Object result = handlerMethod.invoke(environment);
+
+		assertThat(result).isInstanceOf(CompletableFuture.class);
+		CompletableFuture<String> future = (CompletableFuture<String>) result;
+		assertThat(future.get()).isEqualTo("A");
+	}
+
+
+	private static class TestController {
+
+		@Nullable
+		public Callable<String> handleAndReturnCallable() {
+			return () -> "A";
+		}
+
+	}
+
+}

spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/SchemaMappingInvocationTests.java

@@ -16,6 +16,7 @@
 package org.springframework.graphql.data.method.annotation.support;
 
 import java.util.List;
+import java.util.concurrent.Callable;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.atomic.AtomicReference;
 
@@ -28,6 +29,7 @@ import reactor.core.publisher.Mono;
 import reactor.test.StepVerifier;
 
 import org.springframework.context.annotation.AnnotationConfigApplicationContext;
+import org.springframework.core.task.SimpleAsyncTaskExecutor;
 import org.springframework.data.web.ProjectedPayload;
 import org.springframework.graphql.Author;
 import org.springframework.graphql.Book;
@@ -213,8 +215,13 @@ public class SchemaMappingInvocationTests {
 		context.registerBean(BatchLoaderRegistry.class, () -> registry);
 		context.refresh();
 
+		AnnotatedControllerConfigurer configurer = new AnnotatedControllerConfigurer();
+		configurer.setExecutor(new SimpleAsyncTaskExecutor());
+		configurer.setApplicationContext(context);
+		configurer.afterPropertiesSet();
+
 		return GraphQlSetup.schemaResource(BookSource.schema)
-				.runtimeWiringForAnnotatedControllers(context)
+				.runtimeWiring(configurer)
 				.dataLoaders(registry)
 				.toGraphQlService();
 	}
@@ -255,10 +262,12 @@ public class SchemaMappingInvocationTests {
 		}
 
 		@QueryMapping
-		public Author authorById(DataFetchingEnvironment environment, GraphQLContext context) {
-			context.put("key", "value");
-			String id = environment.getArgument("id");
-			return BookSource.getAuthor(Long.parseLong(id));
+		public Callable<Author> authorById(DataFetchingEnvironment environment, GraphQLContext context) {
+			return () -> {
+				context.put("key", "value");
+				String id = environment.getArgument("id");
+				return BookSource.getAuthor(Long.parseLong(id));
+			};
 		}
 
 		@MutationMapping