WebFlux support for handling of early exceptions

This change enables a WebFlux HandlerAdapter to handle not only the
success scenario when a handler is selected, but also any potential
error signal that may occur instead. This makes it possible to
extend ControllerAdvice support to exceptions from handler mapping
such as a 404, 406, 415, and/or even earlier exceptions from the
WebFilter chain.

Closes gh-22991

spring-webflux/src/main/java/org/springframework/web/reactive/DispatchExceptionHandler.java

@@ -0,0 +1,40 @@
+/*
+ * 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.web.reactive;
+
+import reactor.core.publisher.Mono;
+
+import org.springframework.web.server.ServerWebExchange;
+
+/**
+ * Contract to map a {@link Throwable} to a {@link HandlerResult}.
+ *
+ * @author Rossen Stoyanchev
+ * @since 6.0
+ */
+public interface DispatchExceptionHandler {
+
+	/**
+	 * Handler the given exception and resolve it to {@link HandlerResult} that
+	 * can be used for rendering an HTTP response.
+	 * @param exchange the current exchange
+	 * @param ex the exception to handle
+	 * @return a {@code Mono} that emits a {@code HandlerResult} or the original exception
+	 */
+	Mono<HandlerResult> handleError(ServerWebExchange exchange, Throwable ex);
+
+}

spring-webflux/src/main/java/org/springframework/web/reactive/DispatcherHandler.java

@@ -150,8 +150,8 @@ public Mono<Void> handle(ServerWebExchange exchange) {
 				.concatMap(mapping -> mapping.getHandler(exchange))
 				.next()
 				.switchIfEmpty(createNotFoundError())
-				.flatMap(handler -> invokeHandler(exchange, handler))
-				.flatMap(result -> handleResult(exchange, result));
+				.onErrorResume(ex -> handleDispatchError(exchange, ex))
+				.flatMap(handler -> handleRequestWith(exchange, handler));
 	}
 
 	private <R> Mono<R> createNotFoundError() {
@@ -161,14 +161,27 @@ private <R> Mono<R> createNotFoundError() {
 		});
 	}
 
-	private Mono<HandlerResult> invokeHandler(ServerWebExchange exchange, Object handler) {
+	private Mono<Void> handleDispatchError(ServerWebExchange exchange, Throwable ex) {
+		Mono<HandlerResult> resultMono = Mono.error(ex);
+		if (this.handlerAdapters != null) {
+			for (HandlerAdapter adapter : this.handlerAdapters) {
+				if (adapter instanceof DispatchExceptionHandler exceptionHandler) {
+					resultMono = resultMono.onErrorResume(ex2 -> exceptionHandler.handleError(exchange, ex2));
+				}
+			}
+		}
+		return resultMono.flatMap(result -> handleResult(exchange, result));
+	}
+
+	private Mono<Void> handleRequestWith(ServerWebExchange exchange, Object handler) {
 		if (ObjectUtils.nullSafeEquals(exchange.getResponse().getStatusCode(), HttpStatus.FORBIDDEN)) {
 			return Mono.empty();  // CORS rejection
 		}
 		if (this.handlerAdapters != null) {
-			for (HandlerAdapter handlerAdapter : this.handlerAdapters) {
-				if (handlerAdapter.supports(handler)) {
-					return handlerAdapter.handle(exchange, handler);
+			for (HandlerAdapter adapter : this.handlerAdapters) {
+				if (adapter.supports(handler)) {
+					return adapter.handle(exchange, handler)
+							.flatMap(result -> handleResult(exchange, result));
 				}
 			}
 		}
@@ -179,11 +192,10 @@ private Mono<Void> handleResult(ServerWebExchange exchange, HandlerResult result
 		return getResultHandler(result).handleResult(exchange, result)
 				.checkpoint("Handler " + result.getHandler() + " [DispatcherHandler]")
 				.onErrorResume(ex ->
-						result.applyExceptionHandler(ex).flatMap(exResult -> {
-							String text = "Exception handler " + exResult.getHandler() +
-									", error=\"" + ex.getMessage() + "\" [DispatcherHandler]";
-							return getResultHandler(exResult).handleResult(exchange, exResult).checkpoint(text);
-						}));
+						result.applyExceptionHandler(ex).flatMap(exResult ->
+								getResultHandler(exResult).handleResult(exchange, exResult)
+										.checkpoint("Exception handler " + exResult.getHandler() + ", " +
+												"error=\"" + ex.getMessage() + "\" [DispatcherHandler]")));
 	}
 
 	private HandlerResultHandler getResultHandler(HandlerResult handlerResult) {

spring-webflux/src/main/java/org/springframework/web/reactive/HandlerAdapter.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2016 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.
@@ -26,6 +26,14 @@
  * Contract that decouples the {@link DispatcherHandler} from the details of
  * invoking a handler and makes it possible to support any handler type.
  *
+ * <p>A {@code HandlerAdapter} can implement {@link DispatchExceptionHandler}
+ * if it wants to handle an exception that occured before the request is mapped
+ * to a handler. This allows the {@code HandlerAdapter} to expose a consistent
+ * exception handling mechanism for any request handling error.
+ * In Reactive Streams terms, {@link #handle} processes the onNext, while
+ * {@link DispatchExceptionHandler#handleError} processes the onError signal
+ * from the upstream.
+ *
  * @author Rossen Stoyanchev
  * @author Sebastien Deleuze
  * @since 5.0

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2020 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.
@@ -330,38 +330,47 @@ private InvocableHandlerMethod createAttributeMethod(Object bean, Method method)
 	}
 
 	/**
-	 * Find an {@code @ExceptionHandler} method in {@code @ControllerAdvice}
-	 * components or in the controller of the given {@code @RequestMapping} method.
+	 * Look for an {@code @ExceptionHandler} method within the class of the given
+	 * controller method, and also within {@code @ControllerAdvice} classes that
+	 * are applicable to the class of the given controller method.
+	 * @param ex the exception to find a handler for
+	 * @param handlerMethod the controller method that raised the exception, or
+	 * if {@code null}, check only {@code @ControllerAdvice} classes.
 	 */
 	@Nullable
-	public InvocableHandlerMethod getExceptionHandlerMethod(Throwable ex, HandlerMethod handlerMethod) {
-		Class<?> handlerType = handlerMethod.getBeanType();
-
-		// Controller-local first...
-		Object targetBean = handlerMethod.getBean();
-		Method targetMethod = this.exceptionHandlerCache
-				.computeIfAbsent(handlerType, ExceptionHandlerMethodResolver::new)
-				.resolveMethodByThrowable(ex);
+	public InvocableHandlerMethod getExceptionHandlerMethod(Throwable ex, @Nullable HandlerMethod handlerMethod) {
+
+		Class<?> handlerType = (handlerMethod != null ? handlerMethod.getBeanType() : null);
+		Object exceptionHandlerObject = null;
+		Method exceptionHandlerMethod = null;
+
+		if (handlerType != null) {
+			// Controller-local first...
+			exceptionHandlerObject = handlerMethod.getBean();
+			exceptionHandlerMethod = this.exceptionHandlerCache
+					.computeIfAbsent(handlerType, ExceptionHandlerMethodResolver::new)
+					.resolveMethodByThrowable(ex);
+		}
 
-		if (targetMethod == null) {
+		if (exceptionHandlerMethod == null) {
 			// Global exception handlers...
 			for (Map.Entry<ControllerAdviceBean, ExceptionHandlerMethodResolver> entry : this.exceptionHandlerAdviceCache.entrySet()) {
 				ControllerAdviceBean advice = entry.getKey();
 				if (advice.isApplicableToBeanType(handlerType)) {
-					targetBean = advice.resolveBean();
-					targetMethod = entry.getValue().resolveMethodByThrowable(ex);
-					if (targetMethod != null) {
+					exceptionHandlerMethod = entry.getValue().resolveMethodByThrowable(ex);
+					if (exceptionHandlerMethod != null) {
+						exceptionHandlerObject = advice.resolveBean();
 						break;
 					}
 				}
 			}
 		}
 
-		if (targetMethod == null) {
+		if (exceptionHandlerObject == null || exceptionHandlerMethod == null) {
 			return null;
 		}
 
-		InvocableHandlerMethod invocable = new InvocableHandlerMethod(targetBean, targetMethod);
+		InvocableHandlerMethod invocable = new InvocableHandlerMethod(exceptionHandlerObject, exceptionHandlerMethod);
 		invocable.setArgumentResolvers(this.exceptionHandlerResolvers);
 		return invocable;
 	}

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

@@ -38,6 +38,7 @@
 import org.springframework.web.bind.support.WebBindingInitializer;
 import org.springframework.web.method.HandlerMethod;
 import org.springframework.web.reactive.BindingContext;
+import org.springframework.web.reactive.DispatchExceptionHandler;
 import org.springframework.web.reactive.HandlerAdapter;
 import org.springframework.web.reactive.HandlerMapping;
 import org.springframework.web.reactive.HandlerResult;
@@ -52,7 +53,8 @@
  * @author Rossen Stoyanchev
  * @since 5.0
  */
-public class RequestMappingHandlerAdapter implements HandlerAdapter, ApplicationContextAware, InitializingBean {
+public class RequestMappingHandlerAdapter
+		implements HandlerAdapter, DispatchExceptionHandler, ApplicationContextAware, InitializingBean {
 
 	private static final Log logger = LogFactory.getLog(RequestMappingHandlerAdapter.class);
 
@@ -193,7 +195,7 @@ public Mono<HandlerResult> handle(ServerWebExchange exchange, Object handler) {
 		InvocableHandlerMethod invocableMethod = this.methodResolver.getRequestMappingMethod(handlerMethod);
 
 		Function<Throwable, Mono<HandlerResult>> exceptionHandler =
-				ex -> handleException(ex, handlerMethod, bindingContext, exchange);
+				ex -> handleException(exchange, ex, handlerMethod, bindingContext);
 
 		return this.modelInitializer
 				.initModel(handlerMethod, bindingContext, exchange)
@@ -203,23 +205,31 @@ public Mono<HandlerResult> handle(ServerWebExchange exchange, Object handler) {
 				.onErrorResume(exceptionHandler);
 	}
 
-	private Mono<HandlerResult> handleException(Throwable exception, HandlerMethod handlerMethod,
-			BindingContext bindingContext, ServerWebExchange exchange) {
+	private Mono<HandlerResult> handleException(
+			ServerWebExchange exchange, Throwable exception,
+			@Nullable HandlerMethod handlerMethod, @Nullable BindingContext bindingContext) {
 
 		Assert.state(this.methodResolver != null, "Not initialized");
 
 		// Success and error responses may use different content types
 		exchange.getAttributes().remove(HandlerMapping.PRODUCIBLE_MEDIA_TYPES_ATTRIBUTE);
 		exchange.getResponse().getHeaders().clearContentHeaders();
 
-		InvocableHandlerMethod invocable = this.methodResolver.getExceptionHandlerMethod(exception, handlerMethod);
+		InvocableHandlerMethod invocable =
+				this.methodResolver.getExceptionHandlerMethod(exception, handlerMethod);
+
 		if (invocable != null) {
 			ArrayList<Throwable> exceptions = new ArrayList<>();
 			try {
 				if (logger.isDebugEnabled()) {
 					logger.debug(exchange.getLogPrefix() + "Using @ExceptionHandler " + invocable);
 				}
-				bindingContext.getModel().asMap().clear();
+				if (bindingContext != null) {
+					bindingContext.getModel().asMap().clear();
+				}
+				else {
+					bindingContext = new BindingContext();
+				}
 
 				// Expose causes as provided arguments as well
 				Throwable exToExpose = exception;
@@ -245,4 +255,9 @@ private Mono<HandlerResult> handleException(Throwable exception, HandlerMethod h
 		return Mono.error(exception);
 	}
 
+	@Override
+	public Mono<HandlerResult> handleError(ServerWebExchange exchange, Throwable ex) {
+		return handleException(exchange, ex, null, null);
+	}
+
 }