Improve concurrent handling of result in WebAsyncManager

1. Use state transitions
2. Increase synchronized scope in setConcurrentResultAndDispatch

See gh-32342

Author: rstoyanchev

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

@@ -22,6 +22,7 @@
 import java.util.Map;
 import java.util.concurrent.Callable;
 import java.util.concurrent.Future;
+import java.util.concurrent.atomic.AtomicReference;
 
 import javax.servlet.http.HttpServletRequest;
 
@@ -99,12 +100,7 @@ public final class WebAsyncManager {
 	@Nullable
 	private volatile Object[] concurrentResultContext;
 
-	/*
-	 * Whether the concurrentResult is an error. If such errors remain unhandled, some
-	 * Servlet containers will call AsyncListener#onError at the end, after the ASYNC
-	 * and/or the ERROR dispatch (Boot's case), and we need to ignore those.
-	 */
-	private volatile boolean errorHandlingInProgress;
+	private final AtomicReference<State> state = new AtomicReference<>(State.NOT_STARTED);
 
 	private final Map<Object, CallableProcessingInterceptor> callableInterceptors = new LinkedHashMap<>();
 
@@ -257,6 +253,12 @@ public void registerDeferredResultInterceptors(DeferredResultProcessingIntercept
 	 * {@linkplain #getConcurrentResultContext() concurrentResultContext}.
 	 */
 	public void clearConcurrentResult() {
+		if (!this.state.compareAndSet(State.RESULT_SET, State.NOT_STARTED)) {
+			if (logger.isDebugEnabled()) {
+				logger.debug("Unexpected call to clear: [" + this.state.get() + "]");
+			}
+			return;
+		}
 		synchronized (WebAsyncManager.this) {
 			this.concurrentResult = RESULT_NONE;
 			this.concurrentResultContext = null;
@@ -297,6 +299,11 @@ public void startCallableProcessing(final WebAsyncTask<?> webAsyncTask, Object..
 		Assert.notNull(webAsyncTask, "WebAsyncTask must not be null");
 		Assert.state(this.asyncWebRequest != null, "AsyncWebRequest must not be null");
 
+		if (!this.state.compareAndSet(State.NOT_STARTED, State.ASYNC_PROCESSING)) {
+			throw new IllegalStateException(
+					"Unexpected call to startCallableProcessing: [" + this.state.get() + "]");
+		}
+
 		Long timeout = webAsyncTask.getTimeout();
 		if (timeout != null) {
 			this.asyncWebRequest.setTimeout(timeout);
@@ -320,7 +327,7 @@ public void startCallableProcessing(final WebAsyncTask<?> webAsyncTask, Object..
 
 		this.asyncWebRequest.addTimeoutHandler(() -> {
 			if (logger.isDebugEnabled()) {
-				logger.debug("Async request timeout for " + formatUri(this.asyncWebRequest));
+				logger.debug("Servlet container timeout notification for " + formatUri(this.asyncWebRequest));
 			}
 			Object result = interceptorChain.triggerAfterTimeout(this.asyncWebRequest, callable);
 			if (result != CallableProcessingInterceptor.RESULT_NONE) {
@@ -329,14 +336,12 @@ public void startCallableProcessing(final WebAsyncTask<?> webAsyncTask, Object..
 		});
 
 		this.asyncWebRequest.addErrorHandler(ex -> {
-			if (!this.errorHandlingInProgress) {
-				if (logger.isDebugEnabled()) {
-					logger.debug("Async request error for " + formatUri(this.asyncWebRequest) + ": " + ex);
-				}
-				Object result = interceptorChain.triggerAfterError(this.asyncWebRequest, callable, ex);
-				result = (result != CallableProcessingInterceptor.RESULT_NONE ? result : ex);
-				setConcurrentResultAndDispatch(result);
+			if (logger.isDebugEnabled()) {
+				logger.debug("Servlet container error notification for " + formatUri(this.asyncWebRequest) + ": " + ex);
 			}
+			Object result = interceptorChain.triggerAfterError(this.asyncWebRequest, callable, ex);
+			result = (result != CallableProcessingInterceptor.RESULT_NONE ? result : ex);
+			setConcurrentResultAndDispatch(result);
 		});
 
 		this.asyncWebRequest.addCompletionHandler(() ->
@@ -388,33 +393,40 @@ private void logExecutorWarning(AsyncWebRequest asyncWebRequest) {
 	}
 
 	private void setConcurrentResultAndDispatch(@Nullable Object result) {
+		Assert.state(this.asyncWebRequest != null, "AsyncWebRequest must not be null");
 		synchronized (WebAsyncManager.this) {
-			if (this.concurrentResult != RESULT_NONE) {
+			if (!this.state.compareAndSet(State.ASYNC_PROCESSING, State.RESULT_SET)) {
+				if (logger.isDebugEnabled()) {
+					logger.debug("Async result already set: " +
+							"[" + this.state.get() + "], ignored result: " + result +
+							" for " + formatUri(this.asyncWebRequest));
+				}
 				return;
 			}
-			this.concurrentResult = result;
-			this.errorHandlingInProgress = (result instanceof Throwable);
-		}
 
-		Assert.state(this.asyncWebRequest != null, "AsyncWebRequest must not be null");
-		if (this.asyncWebRequest.isAsyncComplete()) {
+			this.concurrentResult = result;
 			if (logger.isDebugEnabled()) {
-				logger.debug("Async result set but request already complete: " + formatUri(this.asyncWebRequest));
+				logger.debug("Async result set to: " + result + " for " + formatUri(this.asyncWebRequest));
+			}
+
+			if (result instanceof Exception) {
+				if (disconnectedClientHelper.checkAndLogClientDisconnectedException((Exception) result)) {
+					return;
+				}
 			}
-			return;
-		}
 
-		if (result instanceof Exception) {
-			if (disconnectedClientHelper.checkAndLogClientDisconnectedException((Exception) result)) {
+			if (this.asyncWebRequest.isAsyncComplete()) {
+				if (logger.isDebugEnabled()) {
+					logger.debug("Async request already completed for " + formatUri(this.asyncWebRequest));
+				}
 				return;
 			}
-		}
 
-		if (logger.isDebugEnabled()) {
-			logger.debug("Async " + (this.errorHandlingInProgress ? "error" : "result set") +
-					", dispatch to " + formatUri(this.asyncWebRequest));
+			if (logger.isDebugEnabled()) {
+				logger.debug("Performing async dispatch for " + formatUri(this.asyncWebRequest));
+			}
+			this.asyncWebRequest.dispatch();
 		}
-		this.asyncWebRequest.dispatch();
 	}
 
 	/**
@@ -437,6 +449,11 @@ public void startDeferredResultProcessing(
 		Assert.notNull(deferredResult, "DeferredResult must not be null");
 		Assert.state(this.asyncWebRequest != null, "AsyncWebRequest must not be null");
 
+		if (!this.state.compareAndSet(State.NOT_STARTED, State.ASYNC_PROCESSING)) {
+			throw new IllegalStateException(
+					"Unexpected call to startDeferredResultProcessing: [" + this.state.get() + "]");
+		}
+
 		Long timeout = deferredResult.getTimeoutValue();
 		if (timeout != null) {
 			this.asyncWebRequest.setTimeout(timeout);
@@ -450,6 +467,9 @@ public void startDeferredResultProcessing(
 		final DeferredResultInterceptorChain interceptorChain = new DeferredResultInterceptorChain(interceptors);
 
 		this.asyncWebRequest.addTimeoutHandler(() -> {
+			if (logger.isDebugEnabled()) {
+				logger.debug("Servlet container timeout notification for " + formatUri(this.asyncWebRequest));
+			}
 			try {
 				interceptorChain.triggerAfterTimeout(this.asyncWebRequest, deferredResult);
 			}
@@ -459,16 +479,17 @@ public void startDeferredResultProcessing(
 		});
 
 		this.asyncWebRequest.addErrorHandler(ex -> {
-			if (!this.errorHandlingInProgress) {
-				try {
-					if (!interceptorChain.triggerAfterError(this.asyncWebRequest, deferredResult, ex)) {
-						return;
-					}
-					deferredResult.setErrorResult(ex);
-				}
-				catch (Throwable interceptorEx) {
-					setConcurrentResultAndDispatch(interceptorEx);
+			if (logger.isDebugEnabled()) {
+				logger.debug("Servlet container error notification for " + formatUri(this.asyncWebRequest));
+			}
+			try {
+				if (!interceptorChain.triggerAfterError(this.asyncWebRequest, deferredResult, ex)) {
+					return;
 				}
+				deferredResult.setErrorResult(ex);
+			}
+			catch (Throwable interceptorEx) {
+				setConcurrentResultAndDispatch(interceptorEx);
 			}
 		});
 
@@ -494,10 +515,13 @@ private void startAsyncProcessing(Object[] processingContext) {
 		synchronized (WebAsyncManager.this) {
 			this.concurrentResult = RESULT_NONE;
 			this.concurrentResultContext = processingContext;
-			this.errorHandlingInProgress = false;
 		}
 
 		Assert.state(this.asyncWebRequest != null, "AsyncWebRequest must not be null");
+		if (logger.isDebugEnabled()) {
+			logger.debug("Started async request for " + formatUri(this.asyncWebRequest));
+		}
+
 		this.asyncWebRequest.startAsync();
 		if (logger.isDebugEnabled()) {
 			logger.debug("Started async request");
@@ -509,4 +533,31 @@ private static String formatUri(AsyncWebRequest asyncWebRequest) {
 		return (request != null ? request.getRequestURI() : "servlet container");
 	}
 
+
+	/**
+	 * Represents a state for {@link WebAsyncManager} to be in.
+	 * <p><pre>
+	 *        NOT_STARTED <------+
+	 *             |             |
+	 *             v             |
+	 *      ASYNC_PROCESSING     |
+	 *             |             |
+	 *             v             |
+	 *         RESULT_SET -------+
+	 * </pre>
+	 * @since 5.3.33
+	 */
+	private enum State {
+
+		/** No async processing in progress. */
+		NOT_STARTED,
+
+		/** Async handling has started, but the result hasn't been set yet. */
+		ASYNC_PROCESSING,
+
+		/** The result is set, and an async dispatch was performed, unless there is a network error. */
+		RESULT_SET
+
+	}
+
 }