Add onError callback to DeferredResult

Issue: SPR-15614

spring-webmvc/src/main/java/org/springframework/web/servlet/mvc/method/annotation/ReactiveTypeHandler.java

@@ -219,6 +219,7 @@ class ReactiveTypeHandler {
 				terminate();
 				this.emitter.complete();
 			});
+			this.emitter.onError(t -> this.emitter.completeWithError(t));
 			subscription.request(1);
 		}
 

spring-webmvc/src/main/java/org/springframework/web/servlet/mvc/method/annotation/ResponseBodyEmitter.java

@@ -19,6 +19,7 @@ package org.springframework.web.servlet.mvc.method.annotation;
 import java.io.IOException;
 import java.util.LinkedHashSet;
 import java.util.Set;
+import java.util.function.Consumer;
 
 import org.springframework.http.MediaType;
 import org.springframework.http.server.ServerHttpResponse;
@@ -74,6 +75,8 @@ public class ResponseBodyEmitter {
 
 	private final DefaultCallback timeoutCallback = new DefaultCallback();
 
+	private final ErrorCallback errorCallback = new ErrorCallback();
+
 	private final DefaultCallback completionCallback = new DefaultCallback();
 
 
@@ -123,6 +126,7 @@ public class ResponseBodyEmitter {
 		}
 		else {
 			this.handler.onTimeout(this.timeoutCallback);
+			this.handler.onError(this.errorCallback);
 			this.handler.onCompletion(this.completionCallback);
 		}
 	}
@@ -168,11 +172,9 @@ public class ResponseBodyEmitter {
 				this.handler.send(object, mediaType);
 			}
 			catch (IOException ex) {
-				completeWithError(ex);
 				throw ex;
 			}
 			catch (Throwable ex) {
-				completeWithError(ex);
 				throw new IllegalStateException("Failed to send " + object, ex);
 			}
 		}
@@ -214,6 +216,15 @@ public class ResponseBodyEmitter {
 		this.timeoutCallback.setDelegate(callback);
 	}
 
+	/**
+	 * Register code to invoke when an error occurred while processing the async request.
+	 * This method is called from a container thread when an error occurred while processing
+	 * an async request.
+	 */
+	public synchronized void onError(Consumer<Throwable> callback) {
+		this.errorCallback.setDelegate(callback);
+	}
+
 	/**
 	 * Register code to invoke when the async request completes. This method is
 	 * called from a container thread when an async request completed for any
@@ -244,6 +255,8 @@ public class ResponseBodyEmitter {
 
 		void onTimeout(Runnable callback);
 
+		void onError(Consumer<Throwable> callback);
+
 		void onCompletion(Runnable callback);
 	}
 
@@ -291,4 +304,22 @@ public class ResponseBodyEmitter {
 		}
 	}
 
+
+	private class ErrorCallback implements Consumer<Throwable> {
+
+		private Consumer<Throwable> delegate;
+
+		public void setDelegate(Consumer<Throwable> callback) {
+			this.delegate = callback;
+		}
+
+		@Override
+		public void accept(Throwable t) {
+			ResponseBodyEmitter.this.complete = true;
+			if (this.delegate != null) {
+				this.delegate.accept(t);
+			}
+		}
+	}
+
 }

spring-webmvc/src/main/java/org/springframework/web/servlet/mvc/method/annotation/ResponseBodyEmitterReturnValueHandler.java

@@ -19,6 +19,8 @@ package org.springframework.web.servlet.mvc.method.annotation;
 import java.io.IOException;
 import java.io.OutputStream;
 import java.util.List;
+import java.util.function.Consumer;
+
 import javax.servlet.ServletRequest;
 import javax.servlet.http.HttpServletResponse;
 
@@ -217,6 +219,11 @@ public class ResponseBodyEmitterReturnValueHandler implements HandlerMethodRetur
 			this.deferredResult.onTimeout(callback);
 		}
 
+		@Override
+		public void onError(Consumer<Throwable> callback) {
+			this.deferredResult.onError(callback);
+		}
+
 		@Override
 		public void onCompletion(Runnable callback) {
 			this.deferredResult.onCompletion(callback);

spring-webmvc/src/test/java/org/springframework/web/servlet/mvc/method/annotation/ReactiveTypeHandlerTests.java

@@ -22,6 +22,7 @@ import java.util.Collections;
 import java.util.List;
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Consumer;
 import java.util.stream.Collectors;
 
 import org.junit.Before;
@@ -380,6 +381,10 @@ public class ReactiveTypeHandlerTests {
 		public void onTimeout(Runnable callback) {
 		}
 
+		@Override
+		public void onError(Consumer<Throwable> callback) {
+		}
+
 		@Override
 		public void onCompletion(Runnable callback) {
 		}

spring-webmvc/src/test/java/org/springframework/web/servlet/mvc/method/annotation/ResponseBodyEmitterReturnValueHandlerTests.java

@@ -19,6 +19,7 @@ import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Consumer;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -175,6 +176,25 @@ public class ResponseBodyEmitterReturnValueHandlerTests {
 		verify(asyncWebRequest).startAsync();
 	}
 
+	@SuppressWarnings("unchecked")
+	@Test
+	public void responseBodyEmitterWithErrorValue() throws Exception {
+
+		AsyncWebRequest asyncWebRequest = mock(AsyncWebRequest.class);
+		WebAsyncUtils.getAsyncManager(this.request).setAsyncWebRequest(asyncWebRequest);
+
+		ResponseBodyEmitter emitter = new ResponseBodyEmitter(19000L);
+		emitter.onError(mock(Consumer.class));
+		emitter.onCompletion(mock(Runnable.class));
+
+		MethodParameter type = on(TestController.class).resolveReturnType(ResponseBodyEmitter.class);
+		this.handler.handleReturnValue(emitter, type, this.mavContainer, this.webRequest);
+
+		verify(asyncWebRequest).addErrorHandler(any(Consumer.class));
+		verify(asyncWebRequest, times(2)).addCompletionHandler(any(Runnable.class));
+		verify(asyncWebRequest).startAsync();
+	}
+
 	@Test
 	public void sseEmitter() throws Exception {
 		MethodParameter type = on(TestController.class).resolveReturnType(SseEmitter.class);

spring-webmvc/src/test/java/org/springframework/web/servlet/mvc/method/annotation/ResponseBodyEmitterTests.java

@@ -102,6 +102,7 @@ public class ResponseBodyEmitterTests {
 	public void sendAfterHandlerInitialized() throws Exception {
 		this.emitter.initialize(this.handler);
 		verify(this.handler).onTimeout(any());
+		verify(this.handler).onError(any());
 		verify(this.handler).onCompletion(any());
 		verifyNoMoreInteractions(this.handler);
 
@@ -119,6 +120,7 @@ public class ResponseBodyEmitterTests {
 	public void sendAfterHandlerInitializedWithError() throws Exception {
 		this.emitter.initialize(this.handler);
 		verify(this.handler).onTimeout(any());
+		verify(this.handler).onError(any());
 		verify(this.handler).onCompletion(any());
 		verifyNoMoreInteractions(this.handler);
 
@@ -137,6 +139,7 @@ public class ResponseBodyEmitterTests {
 	public void sendWithError() throws Exception {
 		this.emitter.initialize(this.handler);
 		verify(this.handler).onTimeout(any());
+		verify(this.handler).onError(any());
 		verify(this.handler).onCompletion(any());
 		verifyNoMoreInteractions(this.handler);
 
@@ -150,7 +153,6 @@ public class ResponseBodyEmitterTests {
 			// expected
 		}
 		verify(this.handler).send("foo", MediaType.TEXT_PLAIN);
-		verify(this.handler).completeWithError(failure);
 		verifyNoMoreInteractions(this.handler);
 	}
 

spring-webmvc/src/test/java/org/springframework/web/servlet/mvc/method/annotation/SseEmitterTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2015 the original author or authors.
+ * Copyright 2002-2017 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.
@@ -18,6 +18,7 @@ package org.springframework.web.servlet.mvc.method.annotation;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.function.Consumer;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -152,6 +153,10 @@ public class SseEmitterTests {
 		public void onTimeout(Runnable callback) {
 		}
 
+		@Override
+		public void onError(Consumer<Throwable> callback) {
+		}
+
 		@Override
 		public void onCompletion(Runnable callback) {
 		}