Controller method async execution on Java 21+

Closes gh-958

Author: rstoyanchev

spring-graphql/src/main/java/org/springframework/graphql/data/federation/EntityHandlerMethod.java

@@ -40,9 +40,9 @@ final class EntityHandlerMethod extends DataFetcherHandlerMethodSupport {
 
 	EntityHandlerMethod(
 			FederationSchemaFactory.EntityMappingInfo info, HandlerMethodArgumentResolverComposite resolvers,
-			@Nullable Executor executor) {
+			@Nullable Executor executor, boolean invokeAsync) {
 
-		super(info.handlerMethod(), resolvers, executor);
+		super(info.handlerMethod(), resolvers, executor, invokeAsync);
 		this.batchHandlerMethod = info.isBatchHandlerMethod();
 	}
 

spring-graphql/src/main/java/org/springframework/graphql/data/federation/FederationSchemaFactory.java

@@ -92,8 +92,8 @@ public void afterPropertiesSet() {
 		super.afterPropertiesSet();
 
 		detectHandlerMethods().forEach((info) ->
-				this.handlerMethods.put(info.typeName(),
-						new EntityHandlerMethod(info, getArgumentResolvers(), getExecutor())));
+				this.handlerMethods.put(info.typeName(), new EntityHandlerMethod(
+						info, getArgumentResolvers(), getExecutor(), shouldInvokeAsync(info.handlerMethod()))));
 
 		if (this.typeResolver == null) {
 			this.typeResolver = new ClassNameTypeResolver();

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

@@ -47,25 +47,43 @@ public abstract class InvocableHandlerMethodSupport extends HandlerMethod {
 	private static final Object NO_VALUE = new Object();
 
 
-	private final boolean hasCallableReturnValue;
-
 	@Nullable
 	private final Executor executor;
 
+	private final boolean hasCallableReturnValue;
+
+	private final boolean invokeAsync;
+
+
 
 	/**
 	 * Create an instance.
 	 * @param handlerMethod the controller method
 	 * @param executor an {@link Executor} to use for {@link Callable} return values
+	 * @deprecated in favor of alternative constructor
 	 */
+	@Deprecated(since = "1.3.0", forRemoval = true)
 	protected InvocableHandlerMethodSupport(HandlerMethod handlerMethod, @Nullable Executor executor) {
+		this(handlerMethod, executor, false);
+	}
+
+	/**
+	 * Create an instance.
+	 * @param handlerMethod the controller method
+	 * @param executor an {@link Executor} to use for {@link Callable} return values
+	 * @param invokeAsync whether to invoke the method through the Executor
+	 * @since 1.3.0
+	 */
+	protected InvocableHandlerMethodSupport(
+			HandlerMethod handlerMethod, @Nullable Executor executor, boolean invokeAsync) {
 		super(handlerMethod.createWithResolvedBean());
 
-		this.hasCallableReturnValue = getReturnType().getParameterType().equals(Callable.class);
 		this.executor = executor;
+		this.hasCallableReturnValue = getReturnType().getParameterType().equals(Callable.class);
+		this.invokeAsync = (invokeAsync && !this.hasCallableReturnValue);
 
-		Assert.isTrue(!this.hasCallableReturnValue || executor != null,
-				"Controller method has Callable return value, but Executor not provided: " +
+		Assert.isTrue((!this.hasCallableReturnValue && !invokeAsync) || executor != null,
+				"Controller method has Callable return value or invokeAsync=true, but Executor not provided: " +
 						handlerMethod.getBridgedMethod().toGenericString());
 	}
 
@@ -81,18 +99,24 @@ protected InvocableHandlerMethodSupport(HandlerMethod handlerMethod, @Nullable E
 	@Nullable
 	protected Object doInvoke(GraphQLContext graphQLContext, Object... argValues) {
 		if (logger.isTraceEnabled()) {
-			logger.trace("Arguments: " + Arrays.toString(argValues));
+			logger.trace("Invoking " + getBridgedMethod().getName() + "(" + Arrays.toString(argValues) + ")");
 		}
 		Method method = getBridgedMethod();
 		try {
 			if (KotlinDetector.isSuspendingFunction(method)) {
 				return invokeSuspendingFunction(getBean(), method, argValues);
 			}
 
-			Object result = method.invoke(getBean(), argValues);
-
-			if (this.hasCallableReturnValue && result != null) {
-				result = adaptCallable(graphQLContext, (Callable<?>) result);
+			Object result;
+			if (this.invokeAsync) {
+				Callable<Object> callable = () -> method.invoke(getBean(), argValues);
+				result = adaptCallable(graphQLContext, callable);
+			}
+			else {
+				result = method.invoke(getBean(), argValues);
+				if (this.hasCallableReturnValue && result != null) {
+					result = adaptCallable(graphQLContext, (Callable<?>) result);
+				}
 			}
 
 			return result;

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

@@ -343,7 +343,8 @@ private void registerDataFetcher(DataFetcherMappingInfo info, RuntimeWiring.Buil
 		DataFetcher<?> dataFetcher;
 		if (!info.isBatchMapping()) {
 			dataFetcher = new SchemaMappingDataFetcher(
-					info, getArgumentResolvers(), this.validationHelper, getExceptionResolver(), getExecutor());
+					info, getArgumentResolvers(), this.validationHelper, getExceptionResolver(),
+					getExecutor(), shouldInvokeAsync(info.getHandlerMethod()));
 		}
 		else {
 			dataFetcher = registerBatchLoader(info);
@@ -366,7 +367,8 @@ private DataFetcher<Object> registerBatchLoader(DataFetcherMappingInfo info) {
 		}
 
 		HandlerMethod handlerMethod = info.getHandlerMethod();
-		BatchLoaderHandlerMethod invocable = new BatchLoaderHandlerMethod(handlerMethod, getExecutor());
+		BatchLoaderHandlerMethod invocable =
+				new BatchLoaderHandlerMethod(handlerMethod, getExecutor(), shouldInvokeAsync(handlerMethod));
 
 		MethodParameter returnType = handlerMethod.getReturnType();
 		Class<?> clazz = returnType.getParameterType();
@@ -441,12 +443,14 @@ static class SchemaMappingDataFetcher implements SelfDescribingDataFetcher<Objec
 		@Nullable
 		private final Executor executor;
 
+		private final boolean invokeAsync;
+
 		private final boolean subscription;
 
 		SchemaMappingDataFetcher(
 				DataFetcherMappingInfo info, HandlerMethodArgumentResolverComposite argumentResolvers,
 				@Nullable ValidationHelper helper, HandlerDataFetcherExceptionResolver exceptionResolver,
-				@Nullable Executor executor) {
+				@Nullable Executor executor, boolean invokeAsync) {
 
 			this.mappingInfo = info;
 			this.argumentResolvers = argumentResolvers;
@@ -457,6 +461,7 @@ static class SchemaMappingDataFetcher implements SelfDescribingDataFetcher<Objec
 			this.exceptionResolver = exceptionResolver;
 
 			this.executor = executor;
+			this.invokeAsync = invokeAsync;
 			this.subscription = this.mappingInfo.getCoordinates().getTypeName().equalsIgnoreCase("Subscription");
 		}
 
@@ -497,7 +502,7 @@ public Object get(DataFetchingEnvironment environment) throws Exception {
 
 			DataFetcherHandlerMethod handlerMethod = new DataFetcherHandlerMethod(
 					getHandlerMethod(), this.argumentResolvers, this.methodValidationHelper,
-					this.executor, this.subscription);
+					this.executor, this.invokeAsync, this.subscription);
 
 			try {
 				Object result = handlerMethod.invoke(environment);

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

@@ -27,6 +27,7 @@
 import java.util.Set;
 import java.util.concurrent.Callable;
 import java.util.concurrent.Executor;
+import java.util.function.Predicate;
 import java.util.stream.Collectors;
 
 import graphql.schema.DataFetcher;
@@ -37,7 +38,9 @@
 import org.springframework.beans.factory.InitializingBean;
 import org.springframework.context.ApplicationContext;
 import org.springframework.context.ApplicationContextAware;
+import org.springframework.core.KotlinDetector;
 import org.springframework.core.MethodIntrospector;
+import org.springframework.core.ReactiveAdapterRegistry;
 import org.springframework.core.annotation.AnnotatedElementUtils;
 import org.springframework.core.convert.ConversionService;
 import org.springframework.format.FormatterRegistrar;
@@ -47,9 +50,11 @@
 import org.springframework.graphql.data.method.HandlerMethodArgumentResolverComposite;
 import org.springframework.graphql.execution.DataFetcherExceptionResolver;
 import org.springframework.lang.Nullable;
+import org.springframework.scheduling.SchedulingTaskExecutor;
 import org.springframework.stereotype.Controller;
 import org.springframework.util.Assert;
 import org.springframework.util.ClassUtils;
+import org.springframework.util.ReflectionUtils;
 
 /**
  * Convenient base for classes that find annotated controller method with argument
@@ -65,6 +70,9 @@ public abstract class AnnotatedControllerDetectionSupport<M> implements Applicat
 			"org.springframework.security.core.context.SecurityContext",
 			AnnotatedControllerDetectionSupport.class.getClassLoader());
 
+	private static final boolean virtualThreadsPresent =
+			(ReflectionUtils.findMethod(Thread.class, "ofVirtual") != null);
+
 	/**
 	 * Bean name prefix for target beans behind scoped proxies. Used to exclude those
 	 * targets from handler method detection, in favor of the corresponding proxies.
@@ -91,6 +99,9 @@ public abstract class AnnotatedControllerDetectionSupport<M> implements Applicat
 	@Nullable
 	private Executor executor;
 
+	private Predicate<HandlerMethod> blockingMethodPredicate =
+			(virtualThreadsPresent) ? new BlockingHandlerMethodPredicate() : ((method) -> false);
+
 	@Nullable
 	private HandlerMethodArgumentResolverComposite argumentResolvers;
 
@@ -148,20 +159,44 @@ public HandlerDataFetcherExceptionResolver getExceptionResolver() {
 
 	/**
 	 * Configure an {@link Executor} to use for asynchronous handling of
-	 * {@link Callable} return values from controller methods.
+	 * {@link Callable} return values from controller methods, as well as for
+	 * {@link #setBlockingMethodPredicate(Predicate) blocking controller methods}
+	 * on Java 21+.
 	 * <p>By default, this is not set in which case controller methods with a
-	 * {@code Callable} return value cannot be registered.
+	 * {@code Callable} return value are not supported, and blocking methods
+	 * will be invoked synchronously.
 	 * @param executor the executor to use
 	 */
 	public void setExecutor(Executor executor) {
 		this.executor = executor;
 	}
 
+	/**
+	 * Return the {@link #setExecutor(Executor) configured Executor}.
+	 */
 	@Nullable
 	public Executor getExecutor() {
 		return this.executor;
 	}
 
+	/**
+	 * Configure a predicate to decide which controller methods are blocking.
+	 * On Java 21+, such methods are invoked asynchronously through the
+	 * {@link #setExecutor(Executor) configured Executor}, unless the executor
+	 * is a thread pool executor as determined via
+	 * {@link SchedulingTaskExecutor#prefersShortLivedTasks() prefersShortLivedTasks}.
+	 * <p>By default, on Java 21+ the predicate returns false for controller
+	 * method return types known to {@link ReactiveAdapterRegistry} as well as
+	 * {@link KotlinDetector#isSuspendingFunction Kotlin suspending functions}.
+	 * On Java 20 and lower, the predicate returns false. You can configure the
+	 * predicate for more control, or alternatively, return {@link Callable}.
+	 * @param predicate the predicate to use
+	 * @since 1.3
+	 */
+	public void setBlockingMethodPredicate(@Nullable Predicate<HandlerMethod> predicate) {
+		this.blockingMethodPredicate = ((predicate != null) ? predicate : (handlerMethod) -> false);
+	}
+
 	/**
 	 * Return the configured argument resolvers.
 	 */
@@ -287,4 +322,20 @@ protected HandlerMethod createHandlerMethod(Method originalMethod, Object handle
 				new HandlerMethod(handler, method);
 	}
 
+	protected boolean shouldInvokeAsync(HandlerMethod handlerMethod) {
+		return (this.blockingMethodPredicate.test(handlerMethod) && this.executor != null &&
+				!(this.executor instanceof SchedulingTaskExecutor ste && ste.prefersShortLivedTasks()));
+	}
+
+
+	private static final class BlockingHandlerMethodPredicate implements Predicate<HandlerMethod> {
+
+		@Override
+		public boolean test(HandlerMethod hm) {
+			Class<?> returnType = hm.getReturnType().getParameterType();
+			return (ReactiveAdapterRegistry.getSharedInstance().getAdapter(returnType) == null &&
+					!KotlinDetector.isSuspendingFunction(hm.getMethod()));
+		}
+	}
+
 }

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

@@ -216,7 +216,7 @@ private Mono<List<GraphQLError>> invokeExceptionHandler(
 
 		DataFetcherHandlerMethod exceptionHandler = new DataFetcherHandlerMethod(
 				new HandlerMethod(controllerOrAdvice, methodHolder.getMethod()), this.argumentResolvers,
-				null, null, false);
+				null, null, false, false);
 
 		List<Throwable> exceptions = new ArrayList<>();
 		try {

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2023 the original author or authors.
+ * Copyright 2002-2024 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.
@@ -20,6 +20,7 @@
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Map;
+import java.util.concurrent.Callable;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Executor;
 import java.util.function.Function;
@@ -59,8 +60,26 @@ public class BatchLoaderHandlerMethod extends InvocableHandlerMethodSupport {
 	private final ParameterNameDiscoverer parameterNameDiscoverer = new DefaultParameterNameDiscoverer();
 
 
+	/**
+	 * Create an instance.
+	 * @param handlerMethod the controller method
+	 * @param executor an {@link Executor} to use for {@link Callable} return values
+	 * @deprecated in favor of alternative constructor
+	 */
+	@Deprecated(since = "1.3.0", forRemoval = true)
 	public BatchLoaderHandlerMethod(HandlerMethod handlerMethod, @Nullable Executor executor) {
-		super(handlerMethod, executor);
+		this(handlerMethod, executor, false);
+	}
+
+	/**
+	 * Create an instance.
+	 * @param handlerMethod the controller method
+	 * @param executor an {@link Executor} to use for {@link Callable} return values
+	 * @param invokeAsync whether to invoke the method through the Executor
+	 * @since 1.3.0
+	 */
+	public BatchLoaderHandlerMethod(HandlerMethod handlerMethod, @Nullable Executor executor, boolean invokeAsync) {
+		super(handlerMethod, executor, invokeAsync);
 	}
 
 

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

@@ -53,13 +53,33 @@ public class DataFetcherHandlerMethod extends DataFetcherHandlerMethodSupport {
 	 * @param validationHelper to apply bean validation with
 	 * @param executor an {@link Executor} to use for {@link Callable} return values
 	 * @param subscription whether the field being fetched is of subscription type
+	 * @deprecated in favor of alternative constructor
 	 */
+	@Deprecated(since = "1.3.0", forRemoval = true)
 	public DataFetcherHandlerMethod(
 			HandlerMethod handlerMethod, HandlerMethodArgumentResolverComposite resolvers,
 			@Nullable BiConsumer<Object, Object[]> validationHelper, @Nullable Executor executor,
 			boolean subscription) {
 
-		super(handlerMethod, resolvers, executor);
+		this(handlerMethod, resolvers, validationHelper, executor, subscription, false);
+	}
+
+	/**
+	 * Constructor with a parent handler method.
+	 * @param handlerMethod the handler method
+	 * @param resolvers the argument resolvers
+	 * @param validationHelper to apply bean validation with
+	 * @param executor an {@link Executor} to use for {@link Callable} return values
+	 * @param subscription whether the field being fetched is of subscription type
+	 * @param invokeAsync whether to invoke the method through the Executor
+	 * @since 1.3.0
+	 */
+	public DataFetcherHandlerMethod(
+			HandlerMethod handlerMethod, HandlerMethodArgumentResolverComposite resolvers,
+			@Nullable BiConsumer<Object, Object[]> validationHelper,
+			@Nullable Executor executor, boolean invokeAsync, boolean subscription) {
+
+		super(handlerMethod, resolvers, executor, invokeAsync);
 		Assert.isTrue(!resolvers.getResolvers().isEmpty(), "No argument resolvers");
 		this.validationHelper = (validationHelper != null) ? validationHelper : (controller, args) -> { };
 		this.subscription = subscription;

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

@@ -48,9 +48,9 @@ public class DataFetcherHandlerMethodSupport extends InvocableHandlerMethodSuppo
 
 	protected DataFetcherHandlerMethodSupport(
 			HandlerMethod handlerMethod, HandlerMethodArgumentResolverComposite resolvers,
-			@Nullable Executor executor) {
+			@Nullable Executor executor, boolean invokeAsync) {
 
-		super(handlerMethod, executor);
+		super(handlerMethod, executor, invokeAsync);
 		this.resolvers = resolvers;
 	}
 

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

@@ -125,7 +125,7 @@ void resolveMono() throws Exception {
 
 		DataFetcherHandlerMethod handlerMethod = new DataFetcherHandlerMethod(
 				new HandlerMethod(new TestController(), TestController.class.getMethod("handleMono", Mono.class)),
-				resolvers, null, null, false);
+				resolvers, null, null, false, false);
 
 		GraphQLContext graphQLContext = new GraphQLContext.Builder().build();
 
@@ -147,7 +147,7 @@ void resolveFromParameterNameWithBatchMapping() throws Exception {
 
 		BatchLoaderHandlerMethod handlerMethod = new BatchLoaderHandlerMethod(
 				new HandlerMethod(controller,
-						TestController.class.getMethod("getAuthors", List.class, Long.class)), null);
+						TestController.class.getMethod("getAuthors", List.class, Long.class)), null, false);
 
 		GraphQLContext context = new GraphQLContext.Builder().build();
 		context.put("id", 123L);