Allow setting timeout in ResponseBodyEmitter

Issue: SPR-13104

spring-web/src/main/java/org/springframework/web/context/request/async/DeferredResult.java

@@ -78,9 +78,12 @@ public class DeferredResult<T> {
 
 	/**
 	 * Create a DeferredResult with a timeout value.
+	 * <p>By default not set in which case the default configured in the MVC
+	 * Java Config or the MVC namespace is used, or if that's not set, then the
+	 * timeout depends on the default of the underlying server.
 	 * @param timeout timeout value in milliseconds
 	 */
-	public DeferredResult(long timeout) {
+	public DeferredResult(Long timeout) {
 		this(timeout, RESULT_NONE);
 	}
 

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

@@ -58,6 +58,8 @@ import org.springframework.util.Assert;
  */
 public class ResponseBodyEmitter {
 
+	private final Long timeout;
+
 	private volatile Handler handler;
 
 	/* Cache for objects sent before handler is set. */
@@ -72,6 +74,33 @@ public class ResponseBodyEmitter {
 	private Runnable completionCallback;
 
 
+	/**
+	 * Create a new ResponseBodyEmitter instance.
+	 */
+	public ResponseBodyEmitter() {
+		this.timeout = null;
+	}
+
+	/**
+	 * Create a ResponseBodyEmitter with a custom timeout value.
+	 * <p>By default not set in which case the default configured in the MVC
+	 * Java Config or the MVC namespace is used, or if that's not set, then the
+	 * timeout depends on the default of the underlying server.
+	 * @param timeout timeout value in milliseconds
+	 */
+	public ResponseBodyEmitter(Long timeout) {
+		this.timeout = timeout;
+	}
+
+
+	/**
+	 * Return the configured timeout value, if any.
+	 */
+	public Long getTimeout() {
+		return this.timeout;
+	}
+
+
 	/**
 	 * Invoked after the response is updated with the status code and headers,
 	 * if the ResponseBodyEmitter is wrapped in a ResponseEntity, but before the

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

@@ -109,7 +109,7 @@ public class ResponseBodyEmitterReturnValueHandler implements HandlerMethodRetur
 		outputMessage.getBody();
 		outputMessage = new StreamingServletServerHttpResponse(outputMessage);
 
-		DeferredResult<?> deferredResult = new DeferredResult<Object>();
+		DeferredResult<?> deferredResult = new DeferredResult<Object>(emitter.getTimeout());
 		WebAsyncUtils.getAsyncManager(webRequest).startDeferredResultProcessing(deferredResult, mavContainer);
 
 		HttpMessageConvertingHandler handler = new HttpMessageConvertingHandler(outputMessage, deferredResult);

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

@@ -15,6 +15,10 @@
  */
 package org.springframework.web.servlet.mvc.method.annotation;
 
+import static org.junit.Assert.*;
+import static org.mockito.Mockito.*;
+import static org.springframework.web.servlet.mvc.method.annotation.SseEmitter.*;
+
 import java.lang.reflect.Method;
 import java.util.Arrays;
 import java.util.List;
@@ -38,13 +42,6 @@ import org.springframework.web.context.request.async.StandardServletAsyncWebRequ
 import org.springframework.web.context.request.async.WebAsyncUtils;
 import org.springframework.web.method.support.ModelAndViewContainer;
 
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertNotNull;
-import static org.junit.Assert.assertNull;
-import static org.junit.Assert.assertTrue;
-import static org.springframework.web.servlet.mvc.method.annotation.SseEmitter.event;
-
 
 /**
  * Unit tests for ResponseBodyEmitterReturnValueHandler.
@@ -126,6 +123,25 @@ public class ResponseBodyEmitterReturnValueHandlerTests {
 		assertNotNull(asyncContext.getDispatchedPath());
 	}
 
+	@Test
+	public void timeoutValueAndCallback() throws Exception {
+
+		AsyncWebRequest asyncWebRequest = mock(AsyncWebRequest.class);
+		WebAsyncUtils.getAsyncManager(this.request).setAsyncWebRequest(asyncWebRequest);
+
+		ResponseBodyEmitter emitter = new ResponseBodyEmitter(19000L);
+		emitter.onTimeout(mock(Runnable.class));
+		emitter.onCompletion(mock(Runnable.class));
+
+		MethodParameter returnType = returnType(TestController.class, "handle");
+		this.handler.handleReturnValue(emitter, returnType, this.mavContainer, this.webRequest);
+
+		verify(asyncWebRequest).setTimeout(19000L);
+		verify(asyncWebRequest).addTimeoutHandler(any(Runnable.class));
+		verify(asyncWebRequest, times(2)).addCompletionHandler(any(Runnable.class));
+		verify(asyncWebRequest).startAsync();
+	}
+
 	@Test
 	public void sseEmitter() throws Exception {
 		MethodParameter returnType = returnType(TestController.class, "handleSse");