From f716c8e9bc8ada04b4060e21680aca793e315cd0 Mon Sep 17 00:00:00 2001
From: Rossen Stoyanchev <rstoyanchev@pivotal.io>
Date: Tue, 4 Apr 2017 16:30:54 -0400
Subject: [PATCH] Consistently check if AsyncContext already completed

Related to SPR-15412
---
 .../reactive/ServletHttpHandlerAdapter.java   | 54 +++++++++----------
 1 file changed, 26 insertions(+), 28 deletions(-)

diff --git a/spring-web/src/main/java/org/springframework/http/server/reactive/ServletHttpHandlerAdapter.java b/spring-web/src/main/java/org/springframework/http/server/reactive/ServletHttpHandlerAdapter.java
index cf6ce52f4ad..7b2463939b8 100644
--- a/spring-web/src/main/java/org/springframework/http/server/reactive/ServletHttpHandlerAdapter.java
+++ b/spring-web/src/main/java/org/springframework/http/server/reactive/ServletHttpHandlerAdapter.java
@@ -108,9 +108,8 @@ public class ServletHttpHandlerAdapter implements Servlet {
 		ServerHttpRequest httpRequest = createRequest(((HttpServletRequest) request), asyncContext);
 		ServerHttpResponse httpResponse = createResponse(((HttpServletResponse) response), asyncContext);
 
-		asyncContext.addListener(TIMEOUT_LISTENER);
-
 		HandlerResultSubscriber subscriber = new HandlerResultSubscriber(asyncContext);
+		asyncContext.addListener(subscriber);
 		this.httpHandler.handle(httpRequest, httpResponse).subscribe(subscriber);
 	}
 
@@ -150,36 +149,12 @@ public class ServletHttpHandlerAdapter implements Servlet {
 	}
 
 
-	private final static AsyncListener TIMEOUT_LISTENER = new AsyncListener() {
-
-		@Override
-		public void onTimeout(AsyncEvent event) throws IOException {
-			event.getAsyncContext().complete();
-		}
-
-		@Override
-		public void onError(AsyncEvent event) throws IOException {
-			event.getAsyncContext().complete();
-		}
-
-		@Override
-		public void onStartAsync(AsyncEvent event) throws IOException {
-			// no-op
-		}
-
-		@Override
-		public void onComplete(AsyncEvent event) throws IOException {
-			// no-op
-		}
-	};
-
-
-	private class HandlerResultSubscriber implements Subscriber<Void> {
+	private class HandlerResultSubscriber implements Subscriber<Void>, AsyncListener {
 
 		private final AsyncContext asyncContext;
 
 
-		public HandlerResultSubscriber(AsyncContext asyncContext) {
+		HandlerResultSubscriber(AsyncContext asyncContext) {
 			this.asyncContext = asyncContext;
 		}
 
@@ -224,6 +199,29 @@ public class ServletHttpHandlerAdapter implements Servlet {
 				// e.g. TIMEOUT_LISTENER (above) may have completed the AsyncContext
 			}
 		}
+
+
+		// AsyncListener...
+
+		@Override
+		public void onTimeout(AsyncEvent event) throws IOException {
+			runIfAsyncNotComplete(() -> event.getAsyncContext().complete());
+		}
+
+		@Override
+		public void onError(AsyncEvent event) throws IOException {
+			runIfAsyncNotComplete(() -> event.getAsyncContext().complete());
+		}
+
+		@Override
+		public void onStartAsync(AsyncEvent event) throws IOException {
+			// No-op
+		}
+
+		@Override
+		public void onComplete(AsyncEvent event) throws IOException {
+			// No-op
+		}
 	}
 
 }