From 7e9437738272a31e1d764972291de42936194ab1 Mon Sep 17 00:00:00 2001
From: Violeta Georgieva <vgeorgieva@pivotal.io>
Date: Thu, 29 Sep 2016 21:04:27 +0300
Subject: [PATCH] Handle async operation events

Problem:
The following exception is observed on an async timeout:
"java.lang.IllegalStateException: It is invalid to call
isReady() when the response has not been put into non-blocking mode"

Current Implementation:
The async operation events sent by the web container are not propagated
to the internal implementation. When timeout/error happens and if the
application does not complete the async operation, the web container
will complete it. At that point if the application tries to read/write,
the operation will fail with an exception (above) that there is not
async operation started.

Proposed Solution:
On async timeout or error, make calls to:
- AbstractRequestBodyPublisher.onError,
- AbstractResponseBodyProcessor.onError,
- AbstractResponseBodyFlushProcessor.onError
As a result of these calls the async operation will be completed and no
more invocations of read/write will be made.
---
 .../AbstractResponseBodyFlushProcessor.java   |  2 +-
 .../reactive/ServletHttpHandlerAdapter.java   | 39 +++++++++++++++++++
 .../reactive/ServletServerHttpRequest.java    |  6 +++
 .../reactive/ServletServerHttpResponse.java   | 17 +++++++-
 4 files changed, 62 insertions(+), 2 deletions(-)

diff --git a/spring-web/src/main/java/org/springframework/http/server/reactive/AbstractResponseBodyFlushProcessor.java b/spring-web/src/main/java/org/springframework/http/server/reactive/AbstractResponseBodyFlushProcessor.java
index fdfc577ef8..57510a9c09 100644
--- a/spring-web/src/main/java/org/springframework/http/server/reactive/AbstractResponseBodyFlushProcessor.java
+++ b/spring-web/src/main/java/org/springframework/http/server/reactive/AbstractResponseBodyFlushProcessor.java
@@ -120,7 +120,7 @@ abstract class AbstractResponseBodyFlushProcessor implements Processor<Publisher
 
 	}
 
-	private void cancel() {
+	protected void cancel() {
 		this.subscription.cancel();
 	}
 
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 a4d0153d87..b965a2f0f3 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
@@ -18,6 +18,8 @@ package org.springframework.http.server.reactive;
 
 import java.io.IOException;
 import javax.servlet.AsyncContext;
+import javax.servlet.AsyncEvent;
+import javax.servlet.AsyncListener;
 import javax.servlet.ServletException;
 import javax.servlet.annotation.WebServlet;
 import javax.servlet.http.HttpServlet;
@@ -89,6 +91,7 @@ public class ServletHttpHandlerAdapter extends HttpServlet {
 				servletRequest, this.dataBufferFactory, this.bufferSize);
 		ServletServerHttpResponse response = new ServletServerHttpResponse(
 				servletResponse, this.dataBufferFactory, this.bufferSize);
+		asyncContext.addListener(new HandlerAsyncEvent(request, response));
 		HandlerResultSubscriber resultSubscriber = new HandlerResultSubscriber(asyncContext);
 		this.handler.handle(request, response).subscribe(resultSubscriber);
 	}
@@ -127,4 +130,40 @@ public class ServletHttpHandlerAdapter extends HttpServlet {
 		}
 	}
 
+	private static final class HandlerAsyncEvent implements AsyncListener {
+		private final ServletServerHttpRequest request;
+		private final ServletServerHttpResponse response;
+
+		public HandlerAsyncEvent(ServletServerHttpRequest request,
+				ServletServerHttpResponse response) {
+			this.request = request;
+			this.response = response;
+		}
+
+		@Override
+		public void onComplete(AsyncEvent event) throws IOException {
+			// no op
+		}
+
+		@Override
+		public void onTimeout(AsyncEvent event) throws IOException {
+			Throwable t = event.getThrowable();
+			if (t == null) {
+				t = new IllegalStateException("Async operation timeout.");
+			}
+			request.onError(t);
+			response.onError(t);
+		}
+
+		@Override
+		public void onError(AsyncEvent event) throws IOException {
+			request.onError(event.getThrowable());
+			response.onError(event.getThrowable());
+		}
+
+		@Override
+		public void onStartAsync(AsyncEvent event) throws IOException {
+			// no op
+		}
+	}
 }
diff --git a/spring-web/src/main/java/org/springframework/http/server/reactive/ServletServerHttpRequest.java b/spring-web/src/main/java/org/springframework/http/server/reactive/ServletServerHttpRequest.java
index 120769829a..65a22d078a 100644
--- a/spring-web/src/main/java/org/springframework/http/server/reactive/ServletServerHttpRequest.java
+++ b/spring-web/src/main/java/org/springframework/http/server/reactive/ServletServerHttpRequest.java
@@ -170,6 +170,12 @@ public class ServletServerHttpRequest extends AbstractServerHttpRequest {
 		}
 	}
 
+	void onError(Throwable t) {
+		if (this.bodyPublisher != null) {
+			this.bodyPublisher.onError(t);
+		}
+	}
+
 	private RequestBodyPublisher createBodyPublisher() throws IOException {
 		RequestBodyPublisher bodyPublisher = new RequestBodyPublisher(
 				this.request.getInputStream(), this.dataBufferFactory, this.bufferSize);
diff --git a/spring-web/src/main/java/org/springframework/http/server/reactive/ServletServerHttpResponse.java b/spring-web/src/main/java/org/springframework/http/server/reactive/ServletServerHttpResponse.java
index 1755d2ab1a..7233f41775 100644
--- a/spring-web/src/main/java/org/springframework/http/server/reactive/ServletServerHttpResponse.java
+++ b/spring-web/src/main/java/org/springframework/http/server/reactive/ServletServerHttpResponse.java
@@ -54,6 +54,8 @@ public class ServletServerHttpResponse extends AbstractListenerServerHttpRespons
 
 	private volatile ResponseBodyProcessor bodyProcessor;
 
+	private volatile ResponseBodyFlushProcessor bodyFlushProcessor;
+
 
 	public ServletServerHttpResponse(HttpServletResponse response,
 			DataBufferFactory dataBufferFactory, int bufferSize) throws IOException {
@@ -116,8 +118,9 @@ public class ServletServerHttpResponse extends AbstractListenerServerHttpRespons
 
 	@Override
 	protected Processor<Publisher<DataBuffer>, Void> createBodyFlushProcessor() {
-		Processor<Publisher<DataBuffer>, Void> processor = new ResponseBodyFlushProcessor();
+		ResponseBodyFlushProcessor processor = new ResponseBodyFlushProcessor();
 		registerListener();
+		bodyFlushProcessor = processor;
 		return processor;
 	}
 
@@ -152,6 +155,18 @@ public class ServletServerHttpResponse extends AbstractListenerServerHttpRespons
 	}
 
 
+	void onError(Throwable t) {
+		if (bodyFlushProcessor != null) {
+			bodyFlushProcessor.cancel();
+			bodyFlushProcessor.onError(t);
+		}
+		if (bodyProcessor != null) {
+			bodyProcessor.cancel();
+			bodyProcessor.onError(t);
+		}
+	}
+
+
 	private class ResponseBodyProcessor extends AbstractResponseBodyProcessor {
 
 		private final ServletOutputStream outputStream;