From 25e7cd577d6855bd56c463ad0ce4ae20bf33d541 Mon Sep 17 00:00:00 2001
From: Violeta Georgieva <vgeorgieva@pivotal.io>
Date: Wed, 19 Oct 2016 10:42:32 +0300
Subject: [PATCH] Handle AsyncListener.onComplete in Servlet adapter

Typically the Mono<Void> from the HttpHandler also reflects the
completion of the request and response body processors and at that
point invoking AsyncContext#complete() from HandlerResultSubscriber
should be sufficient.

This commit explicitly propagates the AsyncListener.onComplete event
to the request and response body processors for added safety.
Technically as mentioned those processors should have completed but
depending on how the controller is written there is a possibility
the body processors may not have completed.

Issue: SPR-14772
---
 .../server/reactive/ServletHttpHandlerAdapter.java   |  9 +++++----
 .../server/reactive/ServletServerHttpRequest.java    |  7 +++++++
 .../server/reactive/ServletServerHttpResponse.java   | 12 ++++++++++++
 3 files changed, 24 insertions(+), 4 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 05e87092b29..4879dc637df 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
@@ -91,7 +91,7 @@ public class ServletHttpHandlerAdapter extends HttpServlet {
 				servletRequest, this.dataBufferFactory, this.bufferSize);
 		ServletServerHttpResponse response = new ServletServerHttpResponse(
 				servletResponse, this.dataBufferFactory, this.bufferSize);
-		asyncContext.addListener(new ErrorHandlingAsyncListener(request, response));
+		asyncContext.addListener(new EventHandlingAsyncListener(request, response));
 		HandlerResultSubscriber resultSubscriber = new HandlerResultSubscriber(asyncContext);
 		this.handler.handle(request, response).subscribe(resultSubscriber);
 	}
@@ -131,14 +131,14 @@ public class ServletHttpHandlerAdapter extends HttpServlet {
 	}
 
 
-	private static final class ErrorHandlingAsyncListener implements AsyncListener {
+	private static final class EventHandlingAsyncListener implements AsyncListener {
 
 		private final ServletServerHttpRequest request;
 
 		private final ServletServerHttpResponse response;
 
 
-		public ErrorHandlingAsyncListener(ServletServerHttpRequest request,
+		public EventHandlingAsyncListener(ServletServerHttpRequest request,
 				ServletServerHttpResponse response) {
 
 			this.request = request;
@@ -169,7 +169,8 @@ public class ServletHttpHandlerAdapter extends HttpServlet {
 
 		@Override
 		public void onComplete(AsyncEvent event) {
-			// no op
+			this.request.handleAsyncListenerComplete();
+			this.response.handleAsyncListenerComplete();
 		}
 	}
 
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 acbf881b103..1b8ae26fa4f 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
@@ -177,6 +177,13 @@ public class ServletServerHttpRequest extends AbstractServerHttpRequest {
 		}
 	}
 
+	/** Handle a complete callback from the Servlet container */
+	void handleAsyncListenerComplete() {
+		if (this.bodyPublisher != null) {
+			this.bodyPublisher.onAllDataRead();
+		}
+	}
+
 	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 2ff1ff99ec8..0abc0ccc515 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
@@ -166,6 +166,18 @@ public class ServletServerHttpResponse extends AbstractListenerServerHttpRespons
 		}
 	}
 
+	/** Handle a complete callback from the Servlet container */
+	void handleAsyncListenerComplete() {
+		if (this.bodyFlushProcessor != null) {
+			this.bodyFlushProcessor.cancel();
+			this.bodyFlushProcessor.onComplete();
+		}
+		if (this.bodyProcessor != null) {
+			this.bodyProcessor.cancel();
+			this.bodyProcessor.onComplete();
+		}
+	}
+
 
 	private class ResponseBodyProcessor extends AbstractResponseBodyProcessor {