From e83fdaa2ddd9f5ace2ae2376b7da2427fc1ecd20 Mon Sep 17 00:00:00 2001
From: Rossen Stoyanchev <rstoyanchev@pivotal.io>
Date: Wed, 5 Apr 2017 12:41:41 -0400
Subject: [PATCH] Revert recent change causing issue in Boot on WildFly

---
 .../reactive/ServletHttpHandlerAdapter.java   | 87 ++++++++++---------
 1 file changed, 47 insertions(+), 40 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 7b2463939b..1be358bef0 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,8 +108,9 @@ public class ServletHttpHandlerAdapter implements Servlet {
 		ServerHttpRequest httpRequest = createRequest(((HttpServletRequest) request), asyncContext);
 		ServerHttpResponse httpResponse = createResponse(((HttpServletResponse) response), asyncContext);
 
+		asyncContext.addListener(TIMEOUT_HANDLER);
+
 		HandlerResultSubscriber subscriber = new HandlerResultSubscriber(asyncContext);
-		asyncContext.addListener(subscriber);
 		this.httpHandler.handle(httpRequest, httpResponse).subscribe(subscriber);
 	}
 
@@ -127,6 +128,23 @@ public class ServletHttpHandlerAdapter implements Servlet {
 				response, context, getDataBufferFactory(), getBufferSize());
 	}
 
+	/**
+	 * We cannot combine TIMEOUT_HANDLER and HandlerResultSubscriber due to:
+	 * https://issues.jboss.org/browse/WFLY-8515
+	 */
+	private static void runIfAsyncNotComplete(AsyncContext asyncContext, Runnable task) {
+		try {
+			if (asyncContext.getRequest().isAsyncStarted()) {
+				task.run();
+			}
+		}
+		catch (IllegalStateException ex) {
+			// Ignore:
+			// AsyncContext recycled and should not be used
+			// e.g. TIMEOUT_LISTENER (above) may have completed the AsyncContext
+		}
+	}
+
 
 	// Other Servlet methods...
 
@@ -149,7 +167,32 @@ public class ServletHttpHandlerAdapter implements Servlet {
 	}
 
 
-	private class HandlerResultSubscriber implements Subscriber<Void>, AsyncListener {
+	private final static AsyncListener TIMEOUT_HANDLER = new AsyncListener() {
+
+		@Override
+		public void onTimeout(AsyncEvent event) throws IOException {
+			AsyncContext context = event.getAsyncContext();
+			runIfAsyncNotComplete(context, context::complete);
+		}
+
+		@Override
+		public void onError(AsyncEvent event) throws IOException {
+			AsyncContext context = event.getAsyncContext();
+			runIfAsyncNotComplete(context, context::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 final AsyncContext asyncContext;
 
@@ -171,7 +214,7 @@ public class ServletHttpHandlerAdapter implements Servlet {
 
 		@Override
 		public void onError(Throwable ex) {
-			runIfAsyncNotComplete(() -> {
+			runIfAsyncNotComplete(this.asyncContext, () -> {
 				logger.error("Could not complete request", ex);
 				HttpServletResponse response = (HttpServletResponse) this.asyncContext.getResponse();
 				response.setStatus(HttpServletResponse.SC_INTERNAL_SERVER_ERROR);
@@ -181,47 +224,11 @@ public class ServletHttpHandlerAdapter implements Servlet {
 
 		@Override
 		public void onComplete() {
-			runIfAsyncNotComplete(() -> {
+			runIfAsyncNotComplete(this.asyncContext, () -> {
 				logger.debug("Successfully completed request");
 				this.asyncContext.complete();
 			});
 		}
-
-		private void runIfAsyncNotComplete(Runnable task) {
-			try {
-				if (this.asyncContext.getRequest().isAsyncStarted()) {
-					task.run();
-				}
-			}
-			catch (IllegalStateException ex) {
-				// Ignore:
-				// AsyncContext recycled and should not be used
-				// 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
-		}
 	}
 
 }