diff --git a/spring-webmvc/src/main/java/org/springframework/web/servlet/mvc/method/annotation/ResponseBodyEmitter.java b/spring-webmvc/src/main/java/org/springframework/web/servlet/mvc/method/annotation/ResponseBodyEmitter.java index a4536a4b9c..7763262686 100644 --- a/spring-webmvc/src/main/java/org/springframework/web/servlet/mvc/method/annotation/ResponseBodyEmitter.java +++ b/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); } } diff --git a/spring-webmvc/src/main/java/org/springframework/web/servlet/mvc/method/annotation/ResponseBodyEmitterReturnValueHandler.java b/spring-webmvc/src/main/java/org/springframework/web/servlet/mvc/method/annotation/ResponseBodyEmitterReturnValueHandler.java index ef7f925d89..fe99556b96 100644 --- a/spring-webmvc/src/main/java/org/springframework/web/servlet/mvc/method/annotation/ResponseBodyEmitterReturnValueHandler.java +++ b/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); + } } diff --git a/spring-webmvc/src/test/java/org/springframework/web/servlet/mvc/method/annotation/ResponseBodyEmitterTests.java b/spring-webmvc/src/test/java/org/springframework/web/servlet/mvc/method/annotation/ResponseBodyEmitterTests.java index 56014307dc..986deca8a0 100644 --- a/spring-webmvc/src/test/java/org/springframework/web/servlet/mvc/method/annotation/ResponseBodyEmitterTests.java +++ b/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); + } + } diff --git a/spring-webmvc/src/test/java/org/springframework/web/servlet/mvc/method/annotation/SseEmitterTests.java b/spring-webmvc/src/test/java/org/springframework/web/servlet/mvc/method/annotation/SseEmitterTests.java index 87ecf7c2e7..494eba5da9 100644 --- a/spring-webmvc/src/test/java/org/springframework/web/servlet/mvc/method/annotation/SseEmitterTests.java +++ b/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) { + } } }