Use DispatchExceptionHandler in HandlerResult

Commit #2878ad added the DispatchExceptionHandler contract for
mapping an error before a handler is selected to a HandlerResult.
The same is also convenient for use in HandlerResult itself which
currently uses a java.util.Function essentially for the same.

See gh-22991

Author: rstoyanchev

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

@@ -189,24 +189,29 @@ private Mono<Void> handleRequestWith(ServerWebExchange exchange, Object handler)
 	}
 
 	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 ->
-								getResultHandler(exResult).handleResult(exchange, exResult)
-										.checkpoint("Exception handler " + exResult.getHandler() + ", " +
-												"error=\"" + ex.getMessage() + "\" [DispatcherHandler]")));
+		Mono<Void> resultMono = doHandleResult(exchange, result, "Handler " + result.getHandler());
+		if (result.getExceptionHandler() != null) {
+			resultMono = resultMono.onErrorResume(ex ->
+					result.getExceptionHandler().handleError(exchange, ex).flatMap(result2 ->
+							doHandleResult(exchange, result2, "Exception handler " +
+									result2.getHandler() + ", error=\"" + ex.getMessage() + "\"")));
+		}
+		return resultMono;
 	}
 
-	private HandlerResultHandler getResultHandler(HandlerResult handlerResult) {
+	private Mono<Void> doHandleResult(
+			ServerWebExchange exchange, HandlerResult handlerResult, String description) {
+
 		if (this.resultHandlers != null) {
 			for (HandlerResultHandler resultHandler : this.resultHandlers) {
 				if (resultHandler.supports(handlerResult)) {
-					return resultHandler;
+					description += " [DispatcherHandler]";
+					return resultHandler.handleResult(exchange, handlerResult).checkpoint(description);
 				}
 			}
 		}
-		throw new IllegalStateException("No HandlerResultHandler for " + handlerResult.getReturnValue());
+		return Mono.error(new IllegalStateException(
+				"No HandlerResultHandler for " + handlerResult.getReturnValue()));
 	}
 
 	@Override

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

@@ -16,8 +16,6 @@
 
 package org.springframework.web.reactive;
 
-import java.util.function.Function;
-
 import reactor.core.publisher.Mono;
 
 import org.springframework.web.server.ServerWebExchange;
@@ -27,7 +25,7 @@
  * 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
+ * if it wants to handle an exception that occurred 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
@@ -54,9 +52,9 @@ public interface HandlerAdapter {
 	 * result that represents an error response.
 	 * <p>Furthermore since an async {@code HandlerResult} may produce an error
 	 * later during result handling implementations are also encouraged to
-	 * {@link HandlerResult#setExceptionHandler(Function) set an exception
-	 * handler} on the {@code HandlerResult} so that may also be applied later
-	 * after result handling.
+	 * {@link HandlerResult#setExceptionHandler(DispatchExceptionHandler)  set
+	 * an exception handler} on the {@code HandlerResult} so that may also be
+	 * applied later after result handling.
 	 * @param exchange current server exchange
 	 * @param handler the selected handler which must have been previously
 	 * checked via {@link #supports(Object)}

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2017 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.
@@ -44,7 +44,10 @@ public class HandlerResult {
 	private final BindingContext bindingContext;
 
 	@Nullable
-	private Function<Throwable, Mono<HandlerResult>> exceptionHandler;
+	private DispatchExceptionHandler exceptionHandler;
+
+	@Nullable
+	private Function<Throwable, Mono<HandlerResult>> exceptionHandlerFunction;
 
 
 	/**
@@ -124,21 +127,50 @@ public Model getModel() {
 		return this.bindingContext.getModel();
 	}
 
+	/**
+	 * A {@link HandlerAdapter} may use this to provide an exception handler
+	 * to use to map exceptions from handling this result into an alternative
+	 * one. Especially when the return value is asynchronous, an exception is
+	 * not be produced at the point of handler invocation, but rather later when
+	 * result handling causes the actual value or an exception to be produced.
+	 * @param exceptionHandler the exception handler to use
+	 * @since 6.0
+	 */
+	public HandlerResult setExceptionHandler(DispatchExceptionHandler exceptionHandler) {
+		this.exceptionHandler = exceptionHandler;
+		return this;
+	}
+
+	/**
+	 * Return the {@link #setExceptionHandler(DispatchExceptionHandler)
+	 * configured} exception handler.
+	 * @since 6.0
+	 */
+	@Nullable
+	public DispatchExceptionHandler getExceptionHandler() {
+		return this.exceptionHandler;
+	}
+
 	/**
 	 * Configure an exception handler that may be used to produce an alternative
 	 * result when result handling fails. Especially for an async return value
 	 * errors may occur after the invocation of the handler.
 	 * @param function the error handler
 	 * @return the current instance
+	 * @deprecated in favor of {@link #setExceptionHandler(DispatchExceptionHandler)}
 	 */
+	@Deprecated(since = "6.0", forRemoval = true)
 	public HandlerResult setExceptionHandler(Function<Throwable, Mono<HandlerResult>> function) {
-		this.exceptionHandler = function;
+		this.exceptionHandler = (exchange, ex) -> function.apply(ex);
+		this.exceptionHandlerFunction = function;
 		return this;
 	}
 
 	/**
 	 * Whether there is an exception handler.
+	 * @deprecated in favor of checking via {@link #getExceptionHandler()}
 	 */
+	@Deprecated(since = "6.0", forRemoval = true)
 	public boolean hasExceptionHandler() {
 		return (this.exceptionHandler != null);
 	}
@@ -147,9 +179,12 @@ public boolean hasExceptionHandler() {
 	 * Apply the exception handler and return the alternative result.
 	 * @param failure the exception
 	 * @return the new result or the same error if there is no exception handler
+	 * @deprecated without a replacement; for internal invocation only, not used as of 6.0
 	 */
+	@Deprecated(since = "6.0", forRemoval = true)
 	public Mono<HandlerResult> applyExceptionHandler(Throwable failure) {
-		return (this.exceptionHandler != null ? this.exceptionHandler.apply(failure) : Mono.error(failure));
+		return (this.exceptionHandlerFunction != null ?
+				this.exceptionHandlerFunction.apply(failure) : Mono.error(failure));
 	}
 
 }

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

@@ -19,7 +19,6 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
-import java.util.function.Function;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
@@ -194,15 +193,15 @@ public Mono<HandlerResult> handle(ServerWebExchange exchange, Object handler) {
 
 		InvocableHandlerMethod invocableMethod = this.methodResolver.getRequestMappingMethod(handlerMethod);
 
-		Function<Throwable, Mono<HandlerResult>> exceptionHandler =
-				ex -> handleException(exchange, ex, handlerMethod, bindingContext);
+		DispatchExceptionHandler exceptionHandler =
+				(exchange2, ex) -> handleException(exchange, ex, handlerMethod, bindingContext);
 
 		return this.modelInitializer
 				.initModel(handlerMethod, bindingContext, exchange)
 				.then(Mono.defer(() -> invocableMethod.invoke(exchange, bindingContext)))
 				.doOnNext(result -> result.setExceptionHandler(exceptionHandler))
 				.doOnNext(result -> bindingContext.saveModel())
-				.onErrorResume(exceptionHandler);
+				.onErrorResume(ex -> exceptionHandler.handleError(exchange, ex));
 	}
 
 	private Mono<HandlerResult> handleException(