Invoke WebMvc.fn error handlers for async errors

This commit makes sure that any error handlers registered on the route
are also applied when an error occurs asynchronously. This commit
applies to asynchronous bodies with both CompletableFuture and Reactive
Streams, as well as completely asynchronous responses.

Closes gh-26831

Author: poutsma

spring-webmvc/src/main/java/org/springframework/web/servlet/function/DefaultAsyncServerResponse.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2020 the original author or authors.
+ * Copyright 2002-2021 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.
@@ -118,7 +118,7 @@ private <R> R delegate(Function<ServerResponse, R> function) {
 	public ModelAndView writeTo(HttpServletRequest request, HttpServletResponse response, Context context)
 			throws ServletException, IOException {
 
-		writeAsync(request, response, createDeferredResult());
+		writeAsync(request, response, createDeferredResult(request));
 		return null;
 	}
 
@@ -140,7 +140,7 @@ static void writeAsync(HttpServletRequest request, HttpServletResponse response,
 
 	}
 
-	private DeferredResult<ServerResponse> createDeferredResult() {
+	private DeferredResult<ServerResponse> createDeferredResult(HttpServletRequest request) {
 		DeferredResult<ServerResponse> result;
 		if (this.timeout != null) {
 			result = new DeferredResult<>(this.timeout.toMillis());
@@ -153,7 +153,13 @@ private DeferredResult<ServerResponse> createDeferredResult() {
 				if (ex instanceof CompletionException && ex.getCause() != null) {
 					ex = ex.getCause();
 				}
-				result.setErrorResult(ex);
+				ServerResponse errorResponse = errorResponse(ex, request);
+				if (errorResponse != null) {
+					result.setResult(errorResponse);
+				}
+				else {
+					result.setErrorResult(ex);
+				}
 			}
 			else {
 				result.setResult(value);

spring-webmvc/src/main/java/org/springframework/web/servlet/function/DefaultEntityResponseBuilder.java

@@ -361,21 +361,27 @@ public CompletionStageEntityResponse(int statusCode, HttpHeaders headers,
 		protected ModelAndView writeToInternal(HttpServletRequest servletRequest, HttpServletResponse servletResponse,
 				Context context) throws ServletException, IOException {
 
-			DeferredResult<?> deferredResult = createDeferredResult(servletRequest, servletResponse, context);
+			DeferredResult<ServerResponse> deferredResult = createDeferredResult(servletRequest, servletResponse, context);
 			DefaultAsyncServerResponse.writeAsync(servletRequest, servletResponse, deferredResult);
 			return null;
 		}
 
-		private DeferredResult<?> createDeferredResult(HttpServletRequest request, HttpServletResponse response,
+		private DeferredResult<ServerResponse> createDeferredResult(HttpServletRequest request, HttpServletResponse response,
 				Context context) {
 
-			DeferredResult<?> result = new DeferredResult<>();
+			DeferredResult<ServerResponse> result = new DeferredResult<>();
 			entity().handle((value, ex) -> {
 				if (ex != null) {
 					if (ex instanceof CompletionException && ex.getCause() != null) {
 						ex = ex.getCause();
 					}
-					result.setErrorResult(ex);
+					ServerResponse errorResponse = errorResponse(ex, request);
+					if (errorResponse != null) {
+						result.setResult(errorResponse);
+					}
+					else {
+						result.setErrorResult(ex);
+					}
 				}
 				else {
 					try {
@@ -468,7 +474,12 @@ public void onNext(T t) {
 
 			@Override
 			public void onError(Throwable t) {
-				this.deferredResult.setErrorResult(t);
+				try {
+					handleError(t, this.servletRequest, this.servletResponse, this.context);
+				}
+				catch (ServletException | IOException handlingThrowable) {
+					this.deferredResult.setErrorResult(handlingThrowable);
+				}
 			}
 
 			@Override

spring-webmvc/src/main/java/org/springframework/web/servlet/function/ErrorHandlingServerResponse.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2020 the original author or authors.
+ * Copyright 2002-2021 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.
@@ -35,7 +35,6 @@
 
 /**
  * Base class for {@link ServerResponse} implementations with error handling.
- *
  * @author Arjen Poutsma
  * @since 5.3
  */
@@ -55,21 +54,36 @@ protected final <T extends ServerResponse> void addErrorHandler(Predicate<Throwa
 	}
 
 	@Nullable
-	protected ModelAndView handleError(Throwable t, HttpServletRequest servletRequest,
+	protected final ModelAndView handleError(Throwable t, HttpServletRequest servletRequest,
 			HttpServletResponse servletResponse, Context context) throws ServletException, IOException {
 
+		ServerResponse serverResponse = errorResponse(t, servletRequest);
+		if (serverResponse != null) {
+			return serverResponse.writeTo(servletRequest, servletResponse, context);
+		}
+		else if (t instanceof ServletException) {
+			throw (ServletException) t;
+		}
+		else if (t instanceof IOException) {
+			throw (IOException) t;
+		}
+		else {
+			throw new ServletException(t);
+		}
+	}
+
+	@Nullable
+	protected final ServerResponse errorResponse(Throwable t, HttpServletRequest servletRequest) {
 		for (ErrorHandler<?> errorHandler : this.errorHandlers) {
 			if (errorHandler.test(t)) {
 				ServerRequest serverRequest = (ServerRequest)
 						servletRequest.getAttribute(RouterFunctions.REQUEST_ATTRIBUTE);
-				ServerResponse serverResponse = errorHandler.handle(t, serverRequest);
-				return serverResponse.writeTo(servletRequest, servletResponse, context);
+				return errorHandler.handle(t, serverRequest);
 			}
 		}
-		throw new ServletException(t);
+		return null;
 	}
 
-
 	private static class ErrorHandler<T extends ServerResponse> {
 
 		private final Predicate<Throwable> predicate;