Add onTimeout() and onCompletion() callbacks to ResponseBodyEmitter

Issue: SPR-12939

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

@@ -67,6 +67,10 @@ public class ResponseBodyEmitter {
 
 	private Throwable failure;
 
+	private Runnable timeoutCallback;
+
+	private Runnable completionCallback;
+
 
 	/**
 	 * Invoked after the response is updated with the status code and headers,
@@ -96,6 +100,12 @@ public class ResponseBodyEmitter {
 					this.handler.complete();
 				}
 			}
+			if (this.timeoutCallback != null) {
+				this.handler.onTimeout(this.timeoutCallback);
+			}
+			if (this.completionCallback != null) {
+				this.handler.onCompletion(this.completionCallback);
+			}
 		}
 	}
 
@@ -179,6 +189,34 @@ public class ResponseBodyEmitter {
 		}
 	}
 
+	/**
+	 * Register code to invoke when the async request times out. This method is
+	 * called from a container thread when an async request times out.
+	 */
+	public void onTimeout(Runnable callback) {
+		synchronized (this) {
+			this.timeoutCallback = callback;
+			if (this.handler != null) {
+				this.handler.onTimeout(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
+	 * reason including timeout and network error. This method is useful for
+	 * detecting that a {@code ResponseBodyEmitter} instance is no longer usable.
+	 */
+	public void onCompletion(Runnable callback) {
+		synchronized (this) {
+			this.completionCallback = callback;
+			if (this.handler != null) {
+				this.handler.onCompletion(callback);
+			}
+		}
+	}
+
 
 	/**
 	 * Handle sent objects and complete request processing.
@@ -190,6 +228,10 @@ public class ResponseBodyEmitter {
 		void complete();
 
 		void completeWithError(Throwable failure);
+
+		void onTimeout(Runnable callback);
+
+		void onCompletion(Runnable callback);
 	}
 
 }

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

@@ -154,6 +154,16 @@ public class ResponseBodyEmitterReturnValueHandler implements HandlerMethodRetur
 		public void completeWithError(Throwable failure) {
 			this.deferredResult.setErrorResult(failure);
 		}
+
+		@Override
+		public void onTimeout(Runnable callback) {
+			this.deferredResult.onTimeout(callback);
+		}
+
+		@Override
+		public void onCompletion(Runnable callback) {
+			this.deferredResult.onCompletion(callback);
+		}
 	}
 
 

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

@@ -20,15 +20,13 @@ import java.io.IOException;
 import org.junit.Before;
 import org.junit.Test;
 import org.mockito.Mock;
+import static org.mockito.Mockito.*;
+import static org.mockito.Mockito.verify;
 import org.mockito.MockitoAnnotations;
 
 import org.springframework.http.MediaType;
 
 import static org.junit.Assert.fail;
-import static org.mockito.Mockito.doThrow;
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.verifyNoMoreInteractions;
-
 
 /**
  * Unit tests for {@link ResponseBodyEmitter}.
@@ -134,4 +132,36 @@ public class ResponseBodyEmitterTests {
 		verifyNoMoreInteractions(this.handler);
 	}
 
+	@Test
+	public void onTimeoutBeforeHandlerInitialized() throws Exception  {
+		Runnable runnable = mock(Runnable.class);
+		this.emitter.onTimeout(runnable);
+		this.emitter.initialize(this.handler);
+		verify(this.handler).onTimeout(runnable);
+	}
+
+	@Test
+	public void onTimeoutAfterHandlerInitialized() throws Exception  {
+		Runnable runnable = mock(Runnable.class);
+		this.emitter.initialize(this.handler);
+		this.emitter.onTimeout(runnable);
+		verify(this.handler).onTimeout(runnable);
+	}
+
+	@Test
+	public void onCompletionBeforeHandlerInitialized() throws Exception  {
+		Runnable runnable = mock(Runnable.class);
+		this.emitter.onCompletion(runnable);
+		this.emitter.initialize(this.handler);
+		verify(this.handler).onCompletion(runnable);
+	}
+
+	@Test
+	public void onCompletionAfterHandlerInitialized() throws Exception  {
+		Runnable runnable = mock(Runnable.class);
+		this.emitter.initialize(this.handler);
+		this.emitter.onCompletion(runnable);
+		verify(this.handler).onCompletion(runnable);
+	}
+
 }

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

@@ -147,6 +147,14 @@ public class SseEmitterTests {
 		@Override
 		public void completeWithError(Throwable failure) {
 		}
+
+		@Override
+		public void onTimeout(Runnable callback) {
+		}
+
+		@Override
+		public void onCompletion(Runnable callback) {
+		}
 	}
 
 }