Refine behavior on error after response committed

If the response is set and we can't change the status through
ServerHttpResponse any more, allow the error signal to propagate and
let the individual server adapters handle it. Ultimately that should
result in closing the connection.

On Servlet containers, we check one last time if the response is
committed (we may not have filled the buffer). If not then save
the exception as a request attribute, dispatch, and re-throw it on the
container thread.

On Undertow access the connection and close it.

On Netty just let the error through to Reactor Netty.

Issue: SPR-16051

Author: rstoyanchev

spring-web/src/main/java/org/springframework/http/server/reactive/ReactorHttpHandlerAdapter.java

@@ -72,12 +72,8 @@ public Mono<Void> apply(HttpServerRequest request, HttpServerResponse response)
 		}
 
 		return this.httpHandler.handle(adaptedRequest, adaptedResponse)
-				.onErrorResume(ex -> {
-					logger.error("Could not complete request", ex);
-					response.status(HttpResponseStatus.INTERNAL_SERVER_ERROR);
-					return Mono.empty();
-				})
-				.doOnSuccess(aVoid -> logger.debug("Successfully completed request"));
+				.doOnError(ex -> logger.error("Handling completed with error", ex))
+				.doOnSuccess(aVoid -> logger.debug("Handling completed with success"));
 	}
 
 }

spring-web/src/main/java/org/springframework/http/server/reactive/ServletHttpHandlerAdapter.java

@@ -18,11 +18,14 @@
 
 import java.io.IOException;
 import java.util.Collection;
+import java.util.concurrent.atomic.AtomicBoolean;
 import javax.servlet.AsyncContext;
 import javax.servlet.AsyncEvent;
 import javax.servlet.AsyncListener;
+import javax.servlet.DispatcherType;
 import javax.servlet.Servlet;
 import javax.servlet.ServletConfig;
+import javax.servlet.ServletException;
 import javax.servlet.ServletRegistration;
 import javax.servlet.ServletRequest;
 import javax.servlet.ServletResponse;
@@ -41,6 +44,7 @@
 import org.springframework.http.HttpMethod;
 import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
+import org.springframework.web.util.NestedServletException;
 
 /**
  * Adapt {@link HttpHandler} to an {@link HttpServlet} using Servlet Async
@@ -56,9 +60,10 @@ public class ServletHttpHandlerAdapter implements Servlet {
 
 	private static final Log logger = LogFactory.getLog(ServletHttpHandlerAdapter.class);
 
-
 	private static final int DEFAULT_BUFFER_SIZE = 8192;
 
+	private static final String WRITE_ERROR_ATTRIBUTE_NAME = ServletHttpHandlerAdapter.class.getName() + ".ERROR";
+
 
 	private final HttpHandler httpHandler;
 
@@ -151,7 +156,14 @@ private String getServletPath(ServletConfig config) {
 
 
 	@Override
-	public void service(ServletRequest request, ServletResponse response) throws IOException {
+	public void service(ServletRequest request, ServletResponse response) throws ServletException, IOException {
+
+		if (DispatcherType.ASYNC.equals(request.getDispatcherType())) {
+			Throwable ex = (Throwable) request.getAttribute(WRITE_ERROR_ATTRIBUTE_NAME);
+			Assert.notNull(ex, "Unexpected async dispatch");
+			throw new NestedServletException("Write publisher error", ex);
+		}
+
 		// Start async before Read/WriteListener registration
 		AsyncContext asyncContext = request.startAsync();
 		asyncContext.setTimeout(-1);
@@ -163,9 +175,11 @@ public void service(ServletRequest request, ServletResponse response) throws IOE
 			httpResponse = new HttpHeadResponseDecorator(httpResponse);
 		}
 
-		asyncContext.addListener(ERROR_LISTENER);
+		AtomicBoolean isCompleted = new AtomicBoolean();
+		HandlerResultAsyncListener listener = new HandlerResultAsyncListener(isCompleted);
+		asyncContext.addListener(listener);
 
-		HandlerResultSubscriber subscriber = new HandlerResultSubscriber(asyncContext);
+		HandlerResultSubscriber subscriber = new HandlerResultSubscriber(asyncContext, isCompleted);
 		this.httpHandler.handle(httpRequest, httpResponse).subscribe(subscriber);
 	}
 
@@ -199,9 +213,9 @@ public void destroy() {
 	 * We cannot combine ERROR_LISTENER and HandlerResultSubscriber due to:
 	 * https://issues.jboss.org/browse/WFLY-8515
 	 */
-	private static void runIfAsyncNotComplete(AsyncContext asyncContext, Runnable task) {
+	private static void runIfAsyncNotComplete(AsyncContext asyncContext, AtomicBoolean isCompleted, Runnable task) {
 		try {
-			if (asyncContext.getRequest().isAsyncStarted()) {
+			if (asyncContext.getRequest().isAsyncStarted() && isCompleted.compareAndSet(false, true)) {
 				task.run();
 			}
 		}
@@ -212,18 +226,27 @@ private static void runIfAsyncNotComplete(AsyncContext asyncContext, Runnable ta
 	}
 
 
-	private final static AsyncListener ERROR_LISTENER = new AsyncListener() {
+	private static class HandlerResultAsyncListener implements AsyncListener {
+
+		private final AtomicBoolean isCompleted;
+
+
+		public HandlerResultAsyncListener(AtomicBoolean isCompleted) {
+			this.isCompleted = isCompleted;
+		}
 
 		@Override
 		public void onTimeout(AsyncEvent event) {
+			logger.debug("Timeout notification from Servlet container");
 			AsyncContext context = event.getAsyncContext();
-			runIfAsyncNotComplete(context, context::complete);
+			runIfAsyncNotComplete(context, this.isCompleted, context::complete);
 		}
 
 		@Override
 		public void onError(AsyncEvent event) {
+			logger.debug("Error notification from Servlet container");
 			AsyncContext context = event.getAsyncContext();
-			runIfAsyncNotComplete(context, context::complete);
+			runIfAsyncNotComplete(context, this.isCompleted, context::complete);
 		}
 
 		@Override
@@ -242,8 +265,12 @@ private class HandlerResultSubscriber implements Subscriber<Void> {
 
 		private final AsyncContext asyncContext;
 
-		public HandlerResultSubscriber(AsyncContext asyncContext) {
+		private final AtomicBoolean isCompleted;
+
+
+		public HandlerResultSubscriber(AsyncContext asyncContext, AtomicBoolean isCompleted) {
 			this.asyncContext = asyncContext;
+			this.isCompleted = isCompleted;
 		}
 
 		@Override
@@ -258,20 +285,30 @@ public void onNext(Void aVoid) {
 
 		@Override
 		public void onError(Throwable ex) {
-			runIfAsyncNotComplete(this.asyncContext, () -> {
-				logger.error("Could not complete request", ex);
-				HttpServletResponse response = (HttpServletResponse) this.asyncContext.getResponse();
-				response.setStatus(HttpServletResponse.SC_INTERNAL_SERVER_ERROR);
-				this.asyncContext.complete();
+			logger.error("Handling completed with error", ex);
+			runIfAsyncNotComplete(this.asyncContext, this.isCompleted, () -> {
+				if (this.asyncContext.getResponse().isCommitted()) {
+					logger.debug("Dispatching into container to raise error");
+					this.asyncContext.getRequest().setAttribute(WRITE_ERROR_ATTRIBUTE_NAME, ex);
+					this.asyncContext.dispatch();
+				}
+				else {
+					try {
+						logger.debug("Setting response status code to 500");
+						this.asyncContext.getResponse().resetBuffer();
+						((HttpServletResponse) this.asyncContext.getResponse()).setStatus(500);
+					}
+					finally {
+						this.asyncContext.complete();
+					}
+				}
 			});
 		}
 
 		@Override
 		public void onComplete() {
-			runIfAsyncNotComplete(this.asyncContext, () -> {
-				logger.debug("Successfully completed request");
-				this.asyncContext.complete();
-			});
+			logger.debug("Handling completed with success");
+			runIfAsyncNotComplete(this.asyncContext, this.isCompleted, this.asyncContext::complete);
 		}
 	}
 

spring-web/src/main/java/org/springframework/http/server/reactive/UndertowHttpHandlerAdapter.java

@@ -16,6 +16,8 @@
 
 package org.springframework.http.server.reactive;
 
+import java.io.IOException;
+
 import io.undertow.server.HttpServerExchange;
 
 import org.apache.commons.logging.Log;
@@ -98,16 +100,26 @@ public void onNext(Void aVoid) {
 
 		@Override
 		public void onError(Throwable ex) {
-			logger.error("Could not complete request", ex);
-			if (!this.exchange.isResponseStarted() && this.exchange.getStatusCode() < 500) {
+			logger.error("Handling completed with error", ex);
+			if (this.exchange.isResponseStarted()) {
+				try {
+					logger.debug("Closing connection");
+					this.exchange.getConnection().close();
+				}
+				catch (IOException e) {
+					// Ignore
+				}
+			}
+			else {
+				logger.debug("Setting response status code to 500");
 				this.exchange.setStatusCode(500);
+				this.exchange.endExchange();
 			}
-			this.exchange.endExchange();
 		}
 
 		@Override
 		public void onComplete() {
-			logger.debug("Successfully completed request");
+			logger.debug("Handling completed with success");
 			this.exchange.endExchange();
 		}
 	}

spring-web/src/main/java/org/springframework/web/server/adapter/HttpWebHandlerAdapter.java

@@ -157,10 +157,7 @@ public LocaleContextResolver getLocaleContextResolver() {
 	public Mono<Void> handle(ServerHttpRequest request, ServerHttpResponse response) {
 		ServerWebExchange exchange = createExchange(request, response);
 		return getDelegate().handle(exchange)
-				.onErrorResume(ex -> {
-					handleFailure(response, ex);
-					return Mono.empty();
-				})
+				.onErrorResume(ex -> handleFailure(response, ex))
 				.then(Mono.defer(response::setComplete));
 	}
 
@@ -169,8 +166,7 @@ protected ServerWebExchange createExchange(ServerHttpRequest request, ServerHttp
 				getCodecConfigurer(), getLocaleContextResolver());
 	}
 
-	private void handleFailure(ServerHttpResponse response, Throwable ex) {
-		boolean statusCodeChanged = response.setStatusCode(HttpStatus.INTERNAL_SERVER_ERROR);
+	private Mono<Void> handleFailure(ServerHttpResponse response, Throwable ex) {
 		if (isDisconnectedClientError(ex)) {
 			if (disconnectedClientLogger.isTraceEnabled()) {
 				disconnectedClientLogger.trace("Looks like the client has gone away", ex);
@@ -180,14 +176,16 @@ else if (disconnectedClientLogger.isDebugEnabled()) {
 						" (For a full stack trace, set the log category '" + DISCONNECTED_CLIENT_LOG_CATEGORY +
 						"' to TRACE level.)");
 			}
+			return Mono.empty();
 		}
-		else if (!statusCodeChanged) {
-			logger.error("Unhandled failure: " + ex.getMessage() + ", " +
-					"response already committed with status=" + response.getStatusCode());
-		}
-		else {
+		if (response.setStatusCode(HttpStatus.INTERNAL_SERVER_ERROR)) {
 			logger.error("Failed to handle request", ex);
+			return Mono.empty();
 		}
+		// After the response is committed, propagate errors to the server..
+		HttpStatus status = response.getStatusCode();
+		logger.error("Unhandled failure: " + ex.getMessage() + ", response already set (status=" + status + ")");
+		return Mono.error(ex);
 	}
 
 	private boolean isDisconnectedClientError(Throwable ex)  {

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

@@ -18,8 +18,10 @@
 
 import java.io.IOException;
 
+import org.junit.Assume;
 import org.junit.Test;
 import org.reactivestreams.Publisher;
+import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
 import org.springframework.context.ApplicationContext;
@@ -28,6 +30,7 @@
 import org.springframework.context.annotation.Configuration;
 import org.springframework.http.HttpHeaders;
 import org.springframework.http.ResponseEntity;
+import org.springframework.http.server.reactive.bootstrap.ReactorHttpServer;
 import org.springframework.web.bind.annotation.ExceptionHandler;
 import org.springframework.web.bind.annotation.GetMapping;
 import org.springframework.web.bind.annotation.RestController;
@@ -53,28 +56,44 @@ protected ApplicationContext initApplicationContext() {
 
 
 	@Test
-	public void controllerThrowingException() throws Exception {
-		String expected = "Recovered from error: State";
-		assertEquals(expected, performGet("/thrown-exception", new HttpHeaders(), String.class).getBody());
+	public void thrownException() throws Exception {
+		doTest("/thrown-exception", "Recovered from error: State");
 	}
 
 	@Test
-	public void controllerThrowingExceptionWithCause() throws Exception {
-		String expected = "Recovered from error: State";
-		assertEquals(expected, performGet("/thrown-exception-with-cause", new HttpHeaders(), String.class).getBody());
+	public void thrownExceptionWithCause() throws Exception {
+		doTest("/thrown-exception-with-cause", "Recovered from error: State");
 	}
 
 	@Test
-	public void controllerThrowingExceptionWithCauseToHandle() throws Exception {
-		String expected = "Recovered from error: IO";
-		String url = "/thrown-exception-with-cause-to-handle";
-		assertEquals(expected, performGet(url, new HttpHeaders(), String.class).getBody());
+	public void thrownExceptionWithCauseToHandle() throws Exception {
+		doTest("/thrown-exception-with-cause-to-handle", "Recovered from error: IO");
 	}
 
 	@Test
-	public void controllerReturnsMonoError() throws Exception {
-		String expected = "Recovered from error: Argument";
-		assertEquals(expected, performGet("/mono-error", new HttpHeaders(), String.class).getBody());
+	public void errorBeforeFirstItem() throws Exception {
+		doTest("/mono-error", "Recovered from error: Argument");
+	}
+
+	@Test // SPR-16051
+	public void exceptionAfterSeveralItems() throws Exception {
+
+		// TODO: uncomment and try after https://github.com/reactor/reactor-netty/issues/231
+		Assume.assumeFalse(server instanceof ReactorHttpServer);
+
+		try {
+			performGet("/SPR-16051", new HttpHeaders(), String.class).getBody();
+			fail();
+		}
+		catch (Throwable ex) {
+			String message = ex.getMessage();
+			assertNotNull(message);
+			assertTrue("Actual: " + message, message.startsWith("Error while extracting response"));
+		}
+	}
+
+	private void doTest(String url, String expected) throws Exception {
+		assertEquals(expected, performGet(url, new HttpHeaders(), String.class).getBody());
 	}
 
 
@@ -110,6 +129,18 @@ public Publisher<String> handleWithError() {
 			return Mono.error(new IllegalArgumentException("Argument"));
 		}
 
+		@GetMapping("/SPR-16051")
+		public Flux<String> errors() {
+			return Flux.range(1, 10000)
+					.map(i -> {
+						if (i == 1000) {
+							throw new RuntimeException("Random error");
+						}
+						return i + ". foo bar";
+					});
+		}
+
+
 		@ExceptionHandler
 		public Publisher<String> handleArgumentException(IOException ex) {
 			return Mono.just("Recovered from error: " + ex.getMessage());

spring-webflux/src/test/java/org/springframework/web/reactive/socket/AbstractWebSocketIntegrationTests.java

@@ -23,7 +23,6 @@
 import java.util.LinkedHashMap;
 import java.util.Map;
 
-import org.apache.tomcat.websocket.WsWebSocketContainer;
 import org.apache.tomcat.websocket.server.WsContextListener;
 import org.junit.After;
 import org.junit.Before;