Use non-blocking thread in WebFlux controller with RequestBody parameter

This commit ensures that `InvocableHandlerMethod` executes the method
on the desired thread if a non-blocking thread is specified, even in the
case where arguments resolution happens on a different thread. This is
notably the case if the method body is resolved as an input argument to
the controller method (`@RequestBody`).

Closes gh-32502

Author: simonbasle

spring-webflux/src/main/java/org/springframework/web/reactive/result/method/InvocableHandlerMethod.java

@@ -36,6 +36,7 @@
 import kotlin.reflect.jvm.KCallablesJvm;
 import kotlin.reflect.jvm.ReflectJvmMapping;
 import reactor.core.publisher.Mono;
+import reactor.core.scheduler.Scheduler;
 
 import org.springframework.core.CoroutinesUtils;
 import org.springframework.core.DefaultParameterNameDiscoverer;
@@ -59,6 +60,12 @@
  * Extension of {@link HandlerMethod} that invokes the underlying method with
  * argument values resolved from the current HTTP request through a list of
  * {@link HandlerMethodArgumentResolver}.
+ * <p>By default, the method invocation happens on the thread from which the
+ * {@code Mono} was subscribed to, or in some cases the thread that emitted one
+ * of the resolved arguments (e.g. when the request body needs to be decoded).
+ * To ensure a predictable thread for the underlying method's invocation,
+ * a {@link Scheduler} can optionally be provided via
+ * {@link #setInvocationScheduler(Scheduler)}.
  *
  * @author Rossen Stoyanchev
  * @author Juergen Hoeller
@@ -85,6 +92,9 @@ public class InvocableHandlerMethod extends HandlerMethod {
 
 	private Class<?>[] validationGroups = EMPTY_GROUPS;
 
+	@Nullable
+	private Scheduler invocationScheduler;
+
 
 	/**
 	 * Create an instance from a {@code HandlerMethod}.
@@ -153,6 +163,13 @@ public void setMethodValidator(@Nullable MethodValidator methodValidator) {
 				methodValidator.determineValidationGroups(getBean(), getBridgedMethod()) : EMPTY_GROUPS);
 	}
 
+	/**
+	 * Set the {@link Scheduler} on which to perform the method invocation.
+	 * @since 6.1.6
+	 */
+	public void setInvocationScheduler(@Nullable Scheduler invocationScheduler) {
+		this.invocationScheduler = invocationScheduler;
+	}
 
 	/**
 	 * Invoke the method for the given exchange.
@@ -165,7 +182,7 @@ public void setMethodValidator(@Nullable MethodValidator methodValidator) {
 	public Mono<HandlerResult> invoke(
 			ServerWebExchange exchange, BindingContext bindingContext, Object... providedArgs) {
 
-		return getMethodArgumentValues(exchange, bindingContext, providedArgs).flatMap(args -> {
+		return getMethodArgumentValuesOnScheduler(exchange, bindingContext, providedArgs).flatMap(args -> {
 			if (shouldValidateArguments() && this.methodValidator != null) {
 				this.methodValidator.applyArgumentValidation(
 						getBean(), getBridgedMethod(), getMethodParameters(), args, this.validationGroups);
@@ -217,14 +234,19 @@ public Mono<HandlerResult> invoke(
 		});
 	}
 
+	private Mono<Object[]> getMethodArgumentValuesOnScheduler(
+			ServerWebExchange exchange, BindingContext bindingContext, Object... providedArgs) {
+		Mono<Object[]> argumentValuesMono = getMethodArgumentValues(exchange, bindingContext, providedArgs);
+		return this.invocationScheduler != null ? argumentValuesMono.publishOn(this.invocationScheduler) : argumentValuesMono;
+	}
+
 	private Mono<Object[]> getMethodArgumentValues(
 			ServerWebExchange exchange, BindingContext bindingContext, Object... providedArgs) {
 
 		MethodParameter[] parameters = getMethodParameters();
 		if (ObjectUtils.isEmpty(parameters)) {
 			return EMPTY_ARGS;
 		}
-
 		List<Mono<Object>> argMonos = new ArrayList<>(parameters.length);
 		for (MethodParameter parameter : parameters) {
 			parameter.initParameterNameDiscovery(this.parameterNameDiscoverer);

spring-webflux/src/main/java/org/springframework/web/reactive/result/method/annotation/ControllerMethodResolver.java

@@ -28,6 +28,7 @@
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
+import reactor.core.scheduler.Scheduler;
 
 import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
 import org.springframework.context.ApplicationContext;
@@ -121,11 +122,19 @@ class ControllerMethodResolver {
 
 	private final Map<Class<?>, SessionAttributesHandler> sessionAttributesHandlerCache = new ConcurrentHashMap<>(64);
 
+	@Nullable
+	private final Scheduler invocationScheduler;
+
+	@Nullable
+	private final Predicate<? super HandlerMethod> blockingMethodPredicate;
+
 
 	ControllerMethodResolver(
 			ArgumentResolverConfigurer customResolvers, ReactiveAdapterRegistry adapterRegistry,
 			ConfigurableApplicationContext context, List<HttpMessageReader<?>> readers,
-			@Nullable WebBindingInitializer webBindingInitializer) {
+			@Nullable WebBindingInitializer webBindingInitializer,
+			@Nullable Scheduler invocationScheduler,
+			@Nullable Predicate<? super HandlerMethod> blockingMethodPredicate) {
 
 		Assert.notNull(customResolvers, "ArgumentResolverConfigurer is required");
 		Assert.notNull(adapterRegistry, "ReactiveAdapterRegistry is required");
@@ -137,6 +146,8 @@ class ControllerMethodResolver {
 		this.requestMappingResolvers = requestMappingResolvers(customResolvers, adapterRegistry, context, readers);
 		this.exceptionHandlerResolvers = exceptionHandlerResolvers(customResolvers, adapterRegistry, context);
 		this.reactiveAdapterRegistry = adapterRegistry;
+		this.invocationScheduler = invocationScheduler;
+		this.blockingMethodPredicate = blockingMethodPredicate;
 
 		if (BEAN_VALIDATION_PRESENT) {
 			this.methodValidator = HandlerMethodValidator.from(webBindingInitializer, null,
@@ -287,6 +298,21 @@ private static Predicate<MethodParameter> methodParamPredicate(
 		};
 	}
 
+	/**
+	 * Return a {@link Scheduler} for the given method if it is considered
+	 * blocking by the underlying blocking method predicate, or null if no
+	 * particular scheduler should be used for this method invocation.
+	 */
+	@Nullable
+	public Scheduler getSchedulerFor(HandlerMethod handlerMethod) {
+		if (this.invocationScheduler != null) {
+			Assert.state(this.blockingMethodPredicate != null, "Expected HandlerMethod Predicate");
+			if (this.blockingMethodPredicate.test(handlerMethod)) {
+				return this.invocationScheduler;
+			}
+		}
+		return null;
+	}
 
 	/**
 	 * Return an {@link InvocableHandlerMethod} for the given
@@ -297,6 +323,9 @@ public InvocableHandlerMethod getRequestMappingMethod(HandlerMethod handlerMetho
 		invocable.setArgumentResolvers(this.requestMappingResolvers);
 		invocable.setReactiveAdapterRegistry(this.reactiveAdapterRegistry);
 		invocable.setMethodValidator(this.methodValidator);
+		//getSchedulerFor returns null if not applicable, which is ok here
+		invocable.setInvocationScheduler(getSchedulerFor(handlerMethod));
+
 		return invocable;
 	}
 

spring-webflux/src/main/java/org/springframework/web/reactive/result/method/annotation/RequestMappingHandlerAdapter.java

@@ -225,7 +225,8 @@ public void afterPropertiesSet() throws Exception {
 
 		this.methodResolver = new ControllerMethodResolver(
 				this.argumentResolverConfigurer, this.reactiveAdapterRegistry, this.applicationContext,
-				this.messageReaders, this.webBindingInitializer);
+				this.messageReaders, this.webBindingInitializer,
+				this.scheduler, this.blockingMethodPredicate);
 
 		this.modelInitializer = new ModelInitializer(this.methodResolver, this.reactiveAdapterRegistry);
 	}
@@ -260,11 +261,9 @@ public Mono<HandlerResult> handle(ServerWebExchange exchange, Object handler) {
 				.doOnNext(result -> result.setExceptionHandler(exceptionHandler))
 				.onErrorResume(ex -> exceptionHandler.handleError(exchange, ex));
 
-		if (this.scheduler != null) {
-			Assert.state(this.blockingMethodPredicate != null, "Expected HandlerMethod Predicate");
-			if (this.blockingMethodPredicate.test(handlerMethod)) {
-				resultMono = resultMono.subscribeOn(this.scheduler);
-			}
+		Scheduler optionalScheduler = this.methodResolver.getSchedulerFor(handlerMethod);
+		if (optionalScheduler != null) {
+			return resultMono.subscribeOn(optionalScheduler);
 		}
 
 		return resultMono;

spring-webflux/src/test/java/org/springframework/web/reactive/result/method/InvocableHandlerMethodTests.java

@@ -26,6 +26,8 @@
 import org.junit.jupiter.api.Test;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
+import reactor.core.scheduler.Scheduler;
+import reactor.core.scheduler.Schedulers;
 import reactor.test.StepVerifier;
 
 import org.springframework.core.io.buffer.DataBuffer;
@@ -75,6 +77,15 @@ void resolveArg() {
 		assertHandlerResultValue(mono, "success:value1");
 	}
 
+	@Test
+	void resolveArgOnSchedulerThread() {
+		this.resolvers.add(stubResolver(Mono.<Object>just("success").publishOn(Schedulers.newSingle("wrong"))));
+		Method method = ResolvableMethod.on(TestController.class).mockCall(o -> o.singleArgThread(null)).method();
+		Mono<HandlerResult> mono = invokeOnScheduler(Schedulers.newSingle("good"), new TestController(), method);
+
+		assertHandlerResultValue(mono, "success on thread: good-", false);
+	}
+
 	@Test
 	void resolveNoArgValue() {
 		this.resolvers.add(stubResolver(Mono.empty()));
@@ -92,6 +103,14 @@ void resolveNoArgs() {
 		assertHandlerResultValue(mono, "success");
 	}
 
+	@Test
+	void resolveNoArgsOnSchedulerThread() {
+		Method method = ResolvableMethod.on(TestController.class).mockCall(o -> o.noArgsThread()).method();
+		Mono<HandlerResult> mono = invokeOnScheduler(Schedulers.newSingle("good"), new TestController(), method);
+
+		assertHandlerResultValue(mono, "on thread: good-", false);
+	}
+
 	@Test
 	void cannotResolveArg() {
 		Method method = ResolvableMethod.on(TestController.class).mockCall(o -> o.singleArg(null)).method();
@@ -229,6 +248,13 @@ private Mono<HandlerResult> invoke(Object handler, Method method, Object... prov
 		return invocable.invoke(this.exchange, new BindingContext(), providedArgs);
 	}
 
+	private Mono<HandlerResult> invokeOnScheduler(Scheduler scheduler, Object handler, Method method, Object... providedArgs) {
+		InvocableHandlerMethod invocable = new InvocableHandlerMethod(handler, method);
+		invocable.setArgumentResolvers(this.resolvers);
+		invocable.setInvocationScheduler(scheduler);
+		return invocable.invoke(this.exchange, new BindingContext(), providedArgs);
+	}
+
 	private HandlerMethodArgumentResolver stubResolver(Object stubValue) {
 		return stubResolver(Mono.just(stubValue));
 	}
@@ -241,8 +267,19 @@ private HandlerMethodArgumentResolver stubResolver(Mono<Object> stubValue) {
 	}
 
 	private void assertHandlerResultValue(Mono<HandlerResult> mono, String expected) {
+		this.assertHandlerResultValue(mono, expected, true);
+	}
+
+	private void assertHandlerResultValue(Mono<HandlerResult> mono, String expected, boolean strict) {
 		StepVerifier.create(mono)
-				.consumeNextWith(result -> assertThat(result.getReturnValue()).isEqualTo(expected))
+				.assertNext(result -> {
+					if (strict) {
+						assertThat(result.getReturnValue()).isEqualTo(expected);
+					}
+					else {
+						assertThat(String.valueOf(result.getReturnValue())).startsWith(expected);
+					}
+				})
 				.expectComplete()
 				.verify();
 	}
@@ -259,6 +296,14 @@ String noArgs() {
 			return "success";
 		}
 
+		String singleArgThread(String q) {
+			return q + " on thread: " + Thread.currentThread().getName();
+		}
+
+		String noArgsThread() {
+			return "on thread: " + Thread.currentThread().getName();
+		}
+
 		void exceptionMethod() {
 			throw new IllegalStateException("boo");
 		}