Refactor async result handling in Spring MVC Test

This change removes the use of a CountDownLatch to wait for the
asynchronously computed controller method return value. Instead we
check in a loop every 200 milliseconds if the result has been set.
If the result is not set within the specified amount of time to wait
an IllegalStateException is raised.

Additional changes:
 - Use AtomicReference to hold the async result
 - Remove @ignore annotations on AsyncTests methods
 - Remove checks for the presence of Servlet 3

Issue: SPR-11516

Author: rstoyanchev

spring-test-mvc/src/main/java/org/springframework/test/web/servlet/DefaultMvcResult.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2012 the original author or authors.
+ * Copyright 2002-2014 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.
@@ -15,13 +15,13 @@
  */
 package org.springframework.test.web.servlet;
 
-import java.util.concurrent.CountDownLatch;
-import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
 
 import javax.servlet.http.HttpServletRequest;
 
 import org.springframework.mock.web.MockHttpServletRequest;
 import org.springframework.mock.web.MockHttpServletResponse;
+import org.springframework.util.Assert;
 import org.springframework.web.servlet.FlashMap;
 import org.springframework.web.servlet.HandlerInterceptor;
 import org.springframework.web.servlet.ModelAndView;
@@ -36,6 +36,9 @@
  */
 class DefaultMvcResult implements MvcResult {
 
+	private static final Object RESULT_NONE = new Object();
+
+
 	private final MockHttpServletRequest mockRequest;
 
 	private final MockHttpServletResponse mockResponse;
@@ -48,9 +51,7 @@ class DefaultMvcResult implements MvcResult {
 
 	private Exception resolvedException;
 
-	private Object asyncResult;
-
-	private CountDownLatch asyncResultLatch;
+	private final AtomicReference<Object> asyncResult = new AtomicReference<Object>(RESULT_NONE);
 
 
 	/**
@@ -106,42 +107,39 @@ public FlashMap getFlashMap() {
 	}
 
 	public void setAsyncResult(Object asyncResult) {
-		this.asyncResult = asyncResult;
+		this.asyncResult.set(asyncResult);
 	}
 
 	public Object getAsyncResult() {
 		return getAsyncResult(-1);
 	}
 
-	public Object getAsyncResult(long timeout) {
+	public Object getAsyncResult(long timeToWait) {
+
 		// MockHttpServletRequest type doesn't have async methods
 		HttpServletRequest request = this.mockRequest;
-		if ((timeout != 0) && request.isAsyncStarted()) {
-			if (timeout == -1) {
-				timeout = request.getAsyncContext().getTimeout();
-			}
-			if (!awaitAsyncResult(timeout)) {
-				throw new IllegalStateException(
-						"Gave up waiting on async result from handler [" + this.handler + "] to complete");
-			}
+		if (request.getAsyncContext() != null) {
+			timeToWait = (timeToWait == -1 ? request.getAsyncContext().getTimeout() : timeToWait);
 		}
-		return this.asyncResult;
-	}
 
-	private boolean awaitAsyncResult(long timeout) {
-		if (this.asyncResultLatch != null) {
-			try {
-				return this.asyncResultLatch.await(timeout, TimeUnit.MILLISECONDS);
-			}
-			catch (InterruptedException e) {
-				return false;
+		if (timeToWait > 0) {
+			long endTime = System.currentTimeMillis() + timeToWait;
+			while (System.currentTimeMillis() < endTime && this.asyncResult.get() == RESULT_NONE) {
+				try {
+					Thread.sleep(200);
+				}
+				catch (InterruptedException ex) {
+					throw new IllegalStateException("Interrupted while waiting for " +
+							"async result to be set for handler [" + this.handler + "]", ex);
+				}
 			}
 		}
-		return true;
-	}
 
-	public void setAsyncResultLatch(CountDownLatch asyncResultLatch) {
-		this.asyncResultLatch = asyncResultLatch;
+		Assert.state(this.asyncResult.get() != RESULT_NONE,
+				"Async result for handler [" + this.handler + "] " +
+						"was not set during the specified timeToWait=" + timeToWait);
+
+		return this.asyncResult.get();
 	}
 
 }

spring-test-mvc/src/main/java/org/springframework/test/web/servlet/MvcResult.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2012 the original author or authors.
+ * Copyright 2002-2014 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.
@@ -76,24 +76,22 @@ public interface MvcResult {
 	FlashMap getFlashMap();
 
 	/**
-	 * Get the result of asynchronous execution or {@code null} if concurrent
-	 * handling did not start. This method will hold and await the completion
-	 * of concurrent handling.
+	 * Get the result of async execution. This method will wait for the async result
+	 * to be set for up to the amount of time configured on the async request,
 	 *
-	 * @throws IllegalStateException if concurrent handling does not complete
-	 * within the allocated async timeout value.
+	 * @throws IllegalStateException if the async result was not set.
 	 */
 	Object getAsyncResult();
 
 	/**
-	 * Get the result of asynchronous execution or {@code null} if concurrent
-	 * handling did not start. This method will wait for up to the given timeout
-	 * for the completion of concurrent handling.
+	 * Get the result of async execution. This method will wait for the async result
+	 * to be set for up to the specified amount of time.
 	 *
-	 * @param timeout how long to wait for the async result to be set in
-	 * milliseconds; if -1, the wait will be as long as the async timeout set
-	 * on the Servlet request
+	 * @param timeToWait how long to wait for the async result to be set, in
+	 * 	milliseconds; if -1, then the async request timeout value is used,
+	 *
+	 * @throws IllegalStateException if the async result was not set.
 	 */
-	Object getAsyncResult(long timeout);
+	Object getAsyncResult(long timeToWait);
 
 }

spring-test-mvc/src/main/java/org/springframework/test/web/servlet/TestDispatcherServlet.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2012 the original author or authors.
+ * Copyright 2002-2014 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.
@@ -58,34 +58,25 @@ public TestDispatcherServlet(WebApplicationContext webApplicationContext) {
 	}
 
 	@Override
-	protected void service(HttpServletRequest request, HttpServletResponse response) throws ServletException, IOException {
-
-		CountDownLatch latch = registerAsyncInterceptors(request);
-		getMvcResult(request).setAsyncResultLatch(latch);
+	protected void service(HttpServletRequest request, HttpServletResponse response)
+			throws ServletException, IOException {
 
+		registerAsyncInterceptors(request);
 		super.service(request, response);
 	}
 
-	private CountDownLatch registerAsyncInterceptors(final HttpServletRequest servletRequest) {
-
-		final CountDownLatch asyncResultLatch = new CountDownLatch(1);
-
+	private void registerAsyncInterceptors(final HttpServletRequest servletRequest) {
 		WebAsyncManager asyncManager = WebAsyncUtils.getAsyncManager(servletRequest);
-
 		asyncManager.registerCallableInterceptor(KEY, new CallableProcessingInterceptorAdapter() {
 			public <T> void postProcess(NativeWebRequest request, Callable<T> task, Object value) throws Exception {
 				getMvcResult(servletRequest).setAsyncResult(value);
-				asyncResultLatch.countDown();
 			}
 		});
 		asyncManager.registerDeferredResultInterceptor(KEY, new DeferredResultProcessingInterceptorAdapter() {
 			public <T> void postProcess(NativeWebRequest request, DeferredResult<T> result, Object value) throws Exception {
 				getMvcResult(servletRequest).setAsyncResult(value);
-				asyncResultLatch.countDown();
 			}
 		});
-
-		return asyncResultLatch;
 	}
 
 	protected DefaultMvcResult getMvcResult(ServletRequest request) {

spring-test-mvc/src/main/java/org/springframework/test/web/servlet/result/PrintingResultHandler.java

@@ -135,7 +135,7 @@ protected void printAsyncResult(MvcResult result) throws Exception {
 		if (servlet3Present) {
 			HttpServletRequest request = result.getRequest();
 			this.printer.printValue("Was async started", request.isAsyncStarted());
-			this.printer.printValue("Async result", result.getAsyncResult(0));
+			this.printer.printValue("Async result", (request.isAsyncStarted() ? result.getAsyncResult(0) : null));
 		}
 	}
 

spring-test-mvc/src/test/java/org/springframework/test/web/servlet/DefaultMvcResultTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2013 the original author or authors.
+ * Copyright 2002-2014 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.
@@ -15,16 +15,23 @@
  */
 package org.springframework.test.web.servlet;
 
-import java.util.concurrent.CountDownLatch;
-import java.util.concurrent.TimeUnit;
-
-import javax.servlet.AsyncContext;
-
 import org.junit.Before;
 import org.junit.Test;
 import org.springframework.mock.web.MockHttpServletRequest;
 
-import static org.mockito.BDDMockito.*;
+import javax.servlet.AsyncContext;
+import javax.servlet.DispatcherType;
+import javax.servlet.ServletException;
+import javax.servlet.ServletRequest;
+import javax.servlet.ServletResponse;
+import javax.servlet.http.Part;
+
+import java.io.IOException;
+import java.util.Collection;
+
+import static org.junit.Assert.assertEquals;
+import static org.mockito.BDDMockito.given;
+import static org.mockito.Mockito.mock;
 
 /**
  * Test fixture for {@link DefaultMvcResult}.
@@ -33,82 +40,73 @@
  */
 public class DefaultMvcResultTests {
 
-	private static final long DEFAULT_TIMEOUT = 10000L;
-
 	private DefaultMvcResult mvcResult;
 
-	private CountDownLatch countDownLatch;
-
-
 	@Before
 	public void setup() {
 		ExtendedMockHttpServletRequest request = new ExtendedMockHttpServletRequest();
-		request.setAsyncStarted(true);
-
-		this.countDownLatch = mock(CountDownLatch.class);
-
 		this.mvcResult = new DefaultMvcResult(request, null);
-		this.mvcResult.setAsyncResultLatch(this.countDownLatch);
-	}
-
-	@Test
-	public void getAsyncResultWithTimeout() throws Exception {
-		long timeout = 1234L;
-		given(this.countDownLatch.await(timeout, TimeUnit.MILLISECONDS)).willReturn(true);
-		this.mvcResult.getAsyncResult(timeout);
-		verify(this.countDownLatch).await(timeout, TimeUnit.MILLISECONDS);
-	}
-
-	@Test
-	public void getAsyncResultWithTimeoutNegativeOne() throws Exception {
-		given(this.countDownLatch.await(DEFAULT_TIMEOUT, TimeUnit.MILLISECONDS)).willReturn(true);
-		this.mvcResult.getAsyncResult(-1);
-		verify(this.countDownLatch).await(DEFAULT_TIMEOUT, TimeUnit.MILLISECONDS);
 	}
 
 	@Test
-	public void getAsyncResultWithoutTimeout() throws Exception {
-		given(this.countDownLatch.await(DEFAULT_TIMEOUT, TimeUnit.MILLISECONDS)).willReturn(true);
-		this.mvcResult.getAsyncResult();
-		verify(this.countDownLatch).await(DEFAULT_TIMEOUT, TimeUnit.MILLISECONDS);
+	public void getAsyncResultSuccess() throws Exception {
+		this.mvcResult.setAsyncResult("Foo");
+		assertEquals("Foo", this.mvcResult.getAsyncResult());
 	}
 
-	@Test
-	public void getAsyncResultWithTimeoutZero() throws Exception {
+	@Test(expected = IllegalStateException.class)
+	public void getAsyncResultFailure() throws Exception {
 		this.mvcResult.getAsyncResult(0);
-		verifyZeroInteractions(this.countDownLatch);
-	}
-
-	@Test(expected=IllegalStateException.class)
-	public void getAsyncResultAndTimeOut() throws Exception {
-		this.mvcResult.getAsyncResult(-1);
-		verify(this.countDownLatch).await(DEFAULT_TIMEOUT, TimeUnit.MILLISECONDS);
 	}
 
 
 	private static class ExtendedMockHttpServletRequest extends MockHttpServletRequest {
 
-		private boolean asyncStarted;
 		private AsyncContext asyncContext;
 
 		public ExtendedMockHttpServletRequest() {
-			super();
 			this.asyncContext = mock(AsyncContext.class);
-			given(this.asyncContext.getTimeout()).willReturn(new Long(DEFAULT_TIMEOUT));
-		}
-
-		public void setAsyncStarted(boolean asyncStarted) {
-			this.asyncStarted = asyncStarted;
+			given(this.asyncContext.getTimeout()).willReturn(0L);
 		}
 
 		@Override
 		public boolean isAsyncStarted() {
-			return this.asyncStarted;
+			return true;
 		}
 
 		@Override
 		public AsyncContext getAsyncContext() {
-			return asyncContext;
+			return this.asyncContext;
+		}
+
+		@Override
+		public Collection<Part> getParts() throws IOException, ServletException {
+			return null;
+		}
+
+		@Override
+		public Part getPart(String name) throws IOException, ServletException {
+			return null;
+		}
+
+		@Override
+		public AsyncContext startAsync() throws IllegalStateException {
+			return this.asyncContext;
+		}
+
+		@Override
+		public AsyncContext startAsync(ServletRequest servletRequest, ServletResponse servletResponse) {
+			return this.asyncContext;
+		}
+
+		@Override
+		public boolean isAsyncSupported() {
+			return true;
+		}
+
+		@Override
+		public DispatcherType getDispatcherType() {
+			return null;
 		}
 	}