Add onTimeout/onCompletion callbacks to DeferredResult

Issue: SPR-9914

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

@@ -37,12 +37,12 @@ public interface AsyncWebRequest extends NativeWebRequest {
 	void setTimeout(Long timeout);
 
 	/**
-	 * Set the handler to use when concurrent handling has timed out.
+	 * Add a handler to invoke when concurrent handling has timed out.
 	 */
-	void setTimeoutHandler(Runnable runnable);
+	void addTimeoutHandler(Runnable runnable);
 
 	/**
-	 * Add a Runnable to be invoked when request processing completes.
+	 * Add a handle to invoke when request processing completes.
 	 */
 	void addCompletionHandler(Runnable runnable);
 

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

@@ -39,8 +39,8 @@ class CallableInterceptorChain {
 	private int preProcessIndex = -1;
 
 
-	public CallableInterceptorChain(Collection<CallableProcessingInterceptor> interceptors) {
-		this.interceptors = new ArrayList<CallableProcessingInterceptor>(interceptors);
+	public CallableInterceptorChain(List<CallableProcessingInterceptor> interceptors) {
+		this.interceptors = interceptors;
 	}
 
 	public void applyPreProcess(NativeWebRequest request, Callable<?> task) throws Exception {

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

@@ -20,6 +20,7 @@ import java.util.concurrent.Callable;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.springframework.util.Assert;
+import org.springframework.web.context.request.NativeWebRequest;
 
 /**
  * {@code DeferredResult} provides an alternative to using a {@link Callable}
@@ -41,6 +42,10 @@ public final class DeferredResult<T> {
 
 	private final Object timeoutResult;
 
+	private Runnable timeoutCallback;
+
+	private Runnable completionCallback;
+
 	private DeferredResultHandler resultHandler;
 
 	private Object result = RESULT_NONE;
@@ -56,7 +61,7 @@ public final class DeferredResult<T> {
 	}
 
 	/**
-	 * Create a DeferredResult with a timeout.
+	 * Create a DeferredResult with a timeout value.
 	 * @param timeout timeout value in milliseconds
 	 */
 	public DeferredResult(long timeout) {
@@ -64,7 +69,8 @@ public final class DeferredResult<T> {
 	}
 
 	/**
-	 * Create a DeferredResult with a timeout and a default result to use on timeout.
+	 * Create a DeferredResult with a timeout value and a default result to use
+	 * in case of timeout.
 	 * @param timeout timeout value in milliseconds; ignored if {@code null}
 	 * @param timeoutResult the result to use
 	 */
@@ -73,13 +79,48 @@ public final class DeferredResult<T> {
 		this.timeout = timeout;
 	}
 
+	/**
+	 * Return {@code true} if this DeferredResult is no longer usable either
+	 * because it was previously set or because the underlying request expired.
+	 * <p>
+	 * The result may have been set with a call to {@link #setResult(Object)},
+	 * or {@link #setErrorResult(Object)}, or as a result of a timeout, if a
+	 * timeout result was provided to the constructor. The request may also
+	 * expire due to a timeout or network error.
+	 */
+	public boolean isSetOrExpired() {
+		return ((this.result != RESULT_NONE) || this.expired);
+	}
+
 	/**
 	 * Return the configured timeout value in milliseconds.
 	 */
-	public Long getTimeoutMilliseconds() {
+	Long getTimeoutValue() {
 		return this.timeout;
 	}
 
+	/**
+	 * Register code to invoke when the async request times out. This method is
+	 * called from a container thread when an async request times out before the
+	 * {@code DeferredResult} has been set. It may invoke
+	 * {@link DeferredResult#setResult(Object) setResult} or
+	 * {@link DeferredResult#setErrorResult(Object) setErrorResult} to resume
+	 * processing.
+	 */
+	public void onTimeout(Runnable callback) {
+		this.timeoutCallback = 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 DeferredResult} instance is no longer usable.
+	 */
+	public void onCompletion(Runnable callback) {
+		this.completionCallback = callback;
+	}
+
 	/**
 	 * Provide a handler to use to handle the result value.
 	 * @param resultHandler the handler
@@ -138,33 +179,29 @@ public final class DeferredResult<T> {
 		return setResultInternal(result);
 	}
 
-	/**
-	 * Return {@code true} if this DeferredResult is no longer usable either
-	 * because it was previously set or because the underlying request expired.
-	 * <p>
-	 * The result may have been set with a call to {@link #setResult(Object)},
-	 * or {@link #setErrorResult(Object)}, or as a result of a timeout, if a
-	 * timeout result was provided to the constructor. The request may also
-	 * expire due to a timeout or network error.
-	 */
-	public boolean isSetOrExpired() {
-		return ((this.result != RESULT_NONE) || this.expired);
-	}
+	DeferredResultProcessingInterceptor getInterceptor() {
+		return new DeferredResultProcessingInterceptorAdapter() {
 
-	/**
-	 * Mark this instance expired so it may no longer be used.
-	 * @return the previous value of the expiration flag
-	 */
-	boolean expire() {
-		synchronized (this) {
-			boolean previous = this.expired;
-			this.expired = true;
-			return previous;
-		}
-	}
+			@Override
+			public <S> void afterTimeout(NativeWebRequest request, DeferredResult<S> deferredResult) {
+				if (timeoutCallback != null) {
+					timeoutCallback.run();
+				}
+				if (DeferredResult.this.timeoutResult != RESULT_NONE) {
+					setResultInternal(timeoutResult);
+				}
+			}
 
-	boolean applyTimeoutResult() {
-		return  (this.timeoutResult != RESULT_NONE) ? setResultInternal(this.timeoutResult) : false;
+			@Override
+			public <S> void afterCompletion(NativeWebRequest request, DeferredResult<S> deferredResult) {
+				synchronized (this) {
+					expired = true;
+				}
+				if (completionCallback != null) {
+					completionCallback.run();
+				}
+			}
+		};
 	}
 
 

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

@@ -15,8 +15,6 @@
  */
 package org.springframework.web.context.request.async;
 
-import java.util.ArrayList;
-import java.util.Collection;
 import java.util.List;
 
 import org.apache.commons.logging.Log;
@@ -38,8 +36,8 @@ class DeferredResultInterceptorChain {
 	private int preProcessingIndex = -1;
 
 
-	public DeferredResultInterceptorChain(Collection<DeferredResultProcessingInterceptor> interceptors) {
-		this.interceptors = new ArrayList<DeferredResultProcessingInterceptor>(interceptors);
+	public DeferredResultInterceptorChain(List<DeferredResultProcessingInterceptor> interceptors) {
+		this.interceptors = interceptors;
 	}
 
 	public void applyPreProcess(NativeWebRequest request, DeferredResult<?> deferredResult) throws Exception {

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

@@ -75,7 +75,7 @@ public interface DeferredResultProcessingInterceptor {
 	 * Invoked from a container thread when an async request times out before
 	 * the {@code DeferredResult} has been set. Implementations may invoke
 	 * {@link DeferredResult#setResult(Object) setResult} or
-	 * {@link DeferredResult#setErrorResult(Object) to resume processing.
+	 * {@link DeferredResult#setErrorResult(Object) setErrorResult} to resume processing.
 	 *
 	 * @param request the current request
 	 * @param deferredResult the DeferredResult for the current request; if the

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

@@ -20,6 +20,7 @@ import java.util.concurrent.Callable;
 import org.springframework.beans.factory.BeanFactory;
 import org.springframework.core.task.AsyncTaskExecutor;
 import org.springframework.util.Assert;
+import org.springframework.web.context.request.NativeWebRequest;
 
 /**
  * Holder for a {@link Callable}, a timeout value, and a task executor.
@@ -27,7 +28,7 @@ import org.springframework.util.Assert;
  * @author Rossen Stoyanchev
  * @since 3.2
  */
-public class AsyncTask<V> {
+public class MvcAsyncTask<V> {
 
 	private final Callable<V> callable;
 
@@ -37,40 +38,51 @@ public class AsyncTask<V> {
 
 	private final AsyncTaskExecutor executor;
 
+	private Callable<V> timeoutCallback;
+
+	private Runnable completionCallback;
+
 	private BeanFactory beanFactory;
 
 
 	/**
-	 * Create an AsyncTask with a timeout value and a Callable.
-	 * @param timeout timeout value in milliseconds
+	 * Create an {@code MvcAsyncTask} wrapping the given {@link Callable}.
 	 * @param callable the callable for concurrent handling
 	 */
-	public AsyncTask(long timeout, Callable<V> callable) {
-		this(timeout, null, null, callable);
-		Assert.notNull(timeout, "Timeout must not be null");
+	public MvcAsyncTask(Callable<V> callable) {
+		this(null, null, null, callable);
 	}
 
 	/**
-	 * Create an AsyncTask with a timeout value, an executor name, and a Callable.
+	 * Create an {@code MvcAsyncTask} with a timeout value and a {@link Callable}.
+	 * @param timeout timeout value in milliseconds
+	 * @param callable the callable for concurrent handling
+	 */
+	public MvcAsyncTask(long timeout, Callable<V> callable) {
+		this(timeout, null, null, callable);
+	}
+
+	/**
+	 * Create an {@code MvcAsyncTask} with a timeout value, an executor name, and a {@link Callable}.
 	 * @param timeout timeout value in milliseconds; ignored if {@code null}
 	 * @param callable the callable for concurrent handling
 	 */
-	public AsyncTask(Long timeout, String executorName, Callable<V> callable) {
+	public MvcAsyncTask(Long timeout, String executorName, Callable<V> callable) {
 		this(timeout, null, executorName, callable);
 		Assert.notNull(executor, "Executor name must not be null");
 	}
 
 	/**
-	 * Create an AsyncTask with a timeout value, an executor instance, and a Callable.
+	 * Create an {@code MvcAsyncTask} with a timeout value, an executor instance, and a Callable.
 	 * @param timeout timeout value in milliseconds; ignored if {@code null}
 	 * @param callable the callable for concurrent handling
 	 */
-	public AsyncTask(Long timeout, AsyncTaskExecutor executor, Callable<V> callable) {
+	public MvcAsyncTask(Long timeout, AsyncTaskExecutor executor, Callable<V> callable) {
 		this(timeout, executor, null, callable);
 		Assert.notNull(executor, "Executor must not be null");
 	}
 
-	private AsyncTask(Long timeout, AsyncTaskExecutor executor, String executorName, Callable<V> callable) {
+	private MvcAsyncTask(Long timeout, AsyncTaskExecutor executor, String executorName, Callable<V> callable) {
 		Assert.notNull(callable, "Callable must not be null");
 		this.callable = callable;
 		this.timeout = timeout;
@@ -111,11 +123,50 @@ public class AsyncTask<V> {
 
 	/**
 	 * A {@link BeanFactory} to use to resolve an executor name. Applications are
-	 * not expected to have to set this property when AsyncTask is used in a
+	 * not expected to have to set this property when {@code MvcAsyncTask} is used in a
 	 * Spring MVC controller.
 	 */
 	public void setBeanFactory(BeanFactory beanFactory) {
 		this.beanFactory = beanFactory;
 	}
 
+
+	/**
+	 * Register code to invoke when the async request times out. This method is
+	 * called from a container thread when an async request times out before the
+	 * {@code Callable} has completed. The callback is executed in the same
+	 * thread and therefore should return without blocking. It may return an
+	 * alternative value to use, including an {@link Exception} or return
+	 * {@link CallableProcessingInterceptor#RESULT_NONE RESULT_NONE}.
+	 */
+	public void onTimeout(Callable<V> callback) {
+		this.timeoutCallback = 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.
+	 */
+	public void onCompletion(Runnable callback) {
+		this.completionCallback = callback;
+	}
+
+	CallableProcessingInterceptor getInterceptor() {
+		return new CallableProcessingInterceptorAdapter() {
+
+			@Override
+			public <T> Object afterTimeout(NativeWebRequest request, Callable<T> task) throws Exception {
+				return (timeoutCallback != null) ? timeoutCallback.call() : CallableProcessingInterceptor.RESULT_NONE;
+			}
+
+			@Override
+			public <T> void afterCompletion(NativeWebRequest request, Callable<T> task) throws Exception {
+				if (completionCallback != null) {
+					completionCallback.run();
+				}
+			}
+		};
+	}
+
 }

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

@@ -41,7 +41,7 @@ public class NoSupportAsyncWebRequest extends ServletWebRequest implements Async
 		// ignored
 	}
 
-	public void setTimeoutHandler(Runnable runnable) {
+	public void addTimeoutHandler(Runnable runnable) {
 		// ignored
 	}
 

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

@@ -49,7 +49,7 @@ public class StandardServletAsyncWebRequest extends ServletWebRequest implements
 
 	private AtomicBoolean asyncCompleted = new AtomicBoolean(false);
 
-	private Runnable timeoutHandler;
+	private final List<Runnable> timeoutHandlers = new ArrayList<Runnable>();
 
 	private final List<Runnable> completionHandlers = new ArrayList<Runnable>();
 
@@ -73,8 +73,8 @@ public class StandardServletAsyncWebRequest extends ServletWebRequest implements
 		this.timeout = timeout;
 	}
 
-	public void setTimeoutHandler(Runnable timeoutHandler) {
-		this.timeoutHandler = timeoutHandler;
+	public void addTimeoutHandler(Runnable timeoutHandler) {
+		this.timeoutHandlers.add(timeoutHandler);
 	}
 
 	public void addCompletionHandler(Runnable runnable) {
@@ -127,8 +127,8 @@ public class StandardServletAsyncWebRequest extends ServletWebRequest implements
 	}
 
 	public void onTimeout(AsyncEvent event) throws IOException {
-		if (this.timeoutHandler != null) {
-			this.timeoutHandler.run();
+		for (Runnable handler : this.timeoutHandlers) {
+			handler.run();
 		}
 	}
 

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

@@ -15,7 +15,9 @@
  */
 package org.springframework.web.context.request.async;
 
+import java.util.ArrayList;
 import java.util.LinkedHashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.concurrent.Callable;
 
@@ -155,19 +157,27 @@ public final class WebAsyncManager {
 		return this.concurrentResultContext;
 	}
 
+	/**
+	 * Get the {@link CallableProcessingInterceptor} registered under the given key.
+	 * @param key the key
+	 * @return the interceptor registered under that key or {@code null}
+	 */
 	public CallableProcessingInterceptor getCallableInterceptor(Object key) {
 		return this.callableInterceptors.get(key);
 	}
 
+	/**
+	 * Get the {@link DeferredResultProcessingInterceptor} registered under the given key.
+	 * @param key the key
+	 * @return the interceptor registered under that key or {@code null}
+	 */
 	public DeferredResultProcessingInterceptor getDeferredResultInterceptor(Object key) {
 		return this.deferredResultInterceptors.get(key);
 	}
 
 	/**
-	 * Register a {@link CallableProcessingInterceptor} that will be applied
-	 * when concurrent request handling with a {@link Callable} starts.
-	 *
-	 * @param key a unique the key under which to register the interceptor
+	 * Register a {@link CallableProcessingInterceptor} under the given key.
+	 * @param key the key
 	 * @param interceptor the interceptor to register
 	 */
 	public void registerCallableInterceptor(Object key, CallableProcessingInterceptor interceptor) {
@@ -181,11 +191,8 @@ public final class WebAsyncManager {
 	}
 
 	/**
-	 * Register a {@link DeferredResultProcessingInterceptor} that will be
-	 * applied when concurrent request handling with a {@link DeferredResult}
-	 * starts.
-	 *
-	 * @param key a unique the key under which to register the interceptor
+	 * Register a {@link DeferredResultProcessingInterceptor} under the given key.
+	 * @param key the key
 	 * @param interceptor the interceptor to register
 	 */
 	public void registerDeferredResultInterceptor(Object key, DeferredResultProcessingInterceptor interceptor) {
@@ -221,16 +228,47 @@ public final class WebAsyncManager {
 	 * @see #getConcurrentResult()
 	 * @see #getConcurrentResultContext()
 	 */
+	@SuppressWarnings({ "rawtypes", "unchecked" })
 	public void startCallableProcessing(final Callable<?> callable, Object... processingContext) {
 		Assert.notNull(callable, "Callable must not be null");
+		startCallableProcessing(new MvcAsyncTask(callable), processingContext);
+	}
+
+	/**
+	 * Use the given {@link MvcAsyncTask} to configure the task executor as well as
+	 * the timeout value of the {@code AsyncWebRequest} before delegating to
+	 * {@link #startCallableProcessing(Callable, Object...)}.
+	 *
+	 * @param mvcAsyncTask an MvcAsyncTask containing the target {@code Callable}
+	 * @param processingContext additional context to save that can be accessed
+	 * via {@link #getConcurrentResultContext()}
+	 */
+	public void startCallableProcessing(final MvcAsyncTask<?> mvcAsyncTask, Object... processingContext) {
+		Assert.notNull(mvcAsyncTask, "MvcAsyncTask must not be null");
 		Assert.state(this.asyncWebRequest != null, "AsyncWebRequest must not be null");
 
-		final CallableInterceptorChain chain = new CallableInterceptorChain(this.callableInterceptors.values());
+		final Callable<?> callable = mvcAsyncTask.getCallable();
 
-		this.asyncWebRequest.setTimeoutHandler(new Runnable() {
+		Long timeout = mvcAsyncTask.getTimeout();
+		if (timeout != null) {
+			this.asyncWebRequest.setTimeout(timeout);
+		}
+
+		AsyncTaskExecutor executor = mvcAsyncTask.getExecutor();
+		if (executor != null) {
+			this.taskExecutor = executor;
+		}
+
+		List<CallableProcessingInterceptor> interceptors = new ArrayList<CallableProcessingInterceptor>();
+		interceptors.add(mvcAsyncTask.getInterceptor());
+		interceptors.addAll(this.callableInterceptors.values());
+
+		final CallableInterceptorChain interceptorChain = new CallableInterceptorChain(interceptors);
+
+		this.asyncWebRequest.addTimeoutHandler(new Runnable() {
 			public void run() {
 				logger.debug("Processing timeout");
-				Object result = chain.triggerAfterTimeout(asyncWebRequest, callable);
+				Object result = interceptorChain.triggerAfterTimeout(asyncWebRequest, callable);
 				if (result != CallableProcessingInterceptor.RESULT_NONE) {
 					setConcurrentResultAndDispatch(result);
 				}
@@ -239,7 +277,7 @@ public final class WebAsyncManager {
 
 		this.asyncWebRequest.addCompletionHandler(new Runnable() {
 			public void run() {
-				chain.triggerAfterCompletion(asyncWebRequest, callable);
+				interceptorChain.triggerAfterCompletion(asyncWebRequest, callable);
 			}
 		});
 
@@ -249,14 +287,14 @@ public final class WebAsyncManager {
 			public void run() {
 				Object result = null;
 				try {
-					chain.applyPreProcess(asyncWebRequest, callable);
+					interceptorChain.applyPreProcess(asyncWebRequest, callable);
 					result = callable.call();
 				}
 				catch (Throwable t) {
 					result = t;
 				}
 				finally {
-					result = chain.applyPostProcess(asyncWebRequest, callable, result);
+					result = interceptorChain.applyPostProcess(asyncWebRequest, callable, result);
 				}
 				setConcurrentResultAndDispatch(result);
 			}
@@ -282,32 +320,6 @@ public final class WebAsyncManager {
 		asyncWebRequest.dispatch();
 	}
 
-	/**
-	 * Use the given {@link AsyncTask} to configure the task executor as well as
-	 * the timeout value of the {@code AsyncWebRequest} before delegating to
-	 * {@link #startCallableProcessing(Callable, Object...)}.
-	 *
-	 * @param asyncTask an asyncTask containing the target {@code Callable}
-	 * @param processingContext additional context to save that can be accessed
-	 * via {@link #getConcurrentResultContext()}
-	 */
-	public void startCallableProcessing(AsyncTask<?> asyncTask, Object... processingContext) {
-		Assert.notNull(asyncTask, "AsyncTask must not be null");
-		Assert.state(this.asyncWebRequest != null, "AsyncWebRequest must not be null");
-
-		Long timeout = asyncTask.getTimeout();
-		if (timeout != null) {
-			this.asyncWebRequest.setTimeout(timeout);
-		}
-
-		AsyncTaskExecutor executor = asyncTask.getExecutor();
-		if (executor != null) {
-			this.taskExecutor = executor;
-		}
-
-		startCallableProcessing(asyncTask.getCallable(), processingContext);
-	}
-
 	/**
 	 * Start concurrent request processing and initialize the given
 	 * {@link DeferredResult} with a {@link DeferredResultHandler} that saves
@@ -329,41 +341,41 @@ public final class WebAsyncManager {
 		Assert.notNull(deferredResult, "DeferredResult must not be null");
 		Assert.state(this.asyncWebRequest != null, "AsyncWebRequest must not be null");
 
-		Long timeout = deferredResult.getTimeoutMilliseconds();
+		Long timeout = deferredResult.getTimeoutValue();
 		if (timeout != null) {
 			this.asyncWebRequest.setTimeout(timeout);
 		}
 
-		final DeferredResultInterceptorChain chain =
-				new DeferredResultInterceptorChain(this.deferredResultInterceptors.values());
+		List<DeferredResultProcessingInterceptor> interceptors = new ArrayList<DeferredResultProcessingInterceptor>();
+		interceptors.add(deferredResult.getInterceptor());
+		interceptors.addAll(this.deferredResultInterceptors.values());
 
-		this.asyncWebRequest.setTimeoutHandler(new Runnable() {
+		final DeferredResultInterceptorChain interceptorChain = new DeferredResultInterceptorChain(interceptors);
+
+		this.asyncWebRequest.addTimeoutHandler(new Runnable() {
 			public void run() {
-				if (!deferredResult.applyTimeoutResult()) {
-					try {
-						chain.triggerAfterTimeout(asyncWebRequest, deferredResult);
-					}
-					catch (Throwable t) {
-						setConcurrentResultAndDispatch(t);
-					}
+				try {
+					interceptorChain.triggerAfterTimeout(asyncWebRequest, deferredResult);
+				}
+				catch (Throwable t) {
+					setConcurrentResultAndDispatch(t);
 				}
 			}
 		});
 
 		this.asyncWebRequest.addCompletionHandler(new Runnable() {
 			public void run() {
-				deferredResult.expire();
-				chain.triggerAfterCompletion(asyncWebRequest, deferredResult);
+				interceptorChain.triggerAfterCompletion(asyncWebRequest, deferredResult);
 			}
 		});
 
 		startAsyncProcessing(processingContext);
 
 		try {
-			chain.applyPreProcess(this.asyncWebRequest, deferredResult);
+			interceptorChain.applyPreProcess(this.asyncWebRequest, deferredResult);
 			deferredResult.setResultHandler(new DeferredResultHandler() {
 				public void handleResult(Object result) {
-					result = chain.applyPostProcess(asyncWebRequest, deferredResult, result);
+					result = interceptorChain.applyPostProcess(asyncWebRequest, deferredResult, result);
 					setConcurrentResultAndDispatch(result);
 				}
 			});