Refactor AbstractResponseBodyFlushProcessor states

With the current state machine
- the implementation can hang after the last element when executing
on Jetty.
- in some cases there will be no flush after the last
Publisher<DataBuffer>.

Author: violetagg

spring-web/src/main/java/org/springframework/http/server/reactive/AbstractResponseBodyFlushProcessor.java

@@ -128,48 +128,63 @@ private enum State {
 			public void onSubscribe(AbstractResponseBodyFlushProcessor processor,
 					Subscription subscription) {
 				Objects.requireNonNull(subscription, "Subscription cannot be null");
-				if (processor.changeState(this, SUBSCRIBED)) {
+				if (processor.changeState(this, REQUESTED)) {
 					processor.subscription = subscription;
 					subscription.request(1);
 				}
 				else {
 					super.onSubscribe(processor, subscription);
 				}
 			}
-		}, SUBSCRIBED {
+		},
+		REQUESTED {
 			@Override
 			public void onNext(AbstractResponseBodyFlushProcessor processor,
 					Publisher<DataBuffer> chunk) {
-				Processor<DataBuffer, Void> chunkProcessor =
-						processor.createBodyProcessor();
-				chunk.subscribe(chunkProcessor);
-				chunkProcessor.subscribe(new WriteSubscriber(processor));
+				if (processor.changeState(this, RECEIVED)) {
+					Processor<DataBuffer, Void> chunkProcessor =
+							processor.createBodyProcessor();
+					chunk.subscribe(chunkProcessor);
+					chunkProcessor.subscribe(new WriteSubscriber(processor));
+				}
 			}
 
 			@Override
 			void onComplete(AbstractResponseBodyFlushProcessor processor) {
-				processor.subscriberCompleted = true;
+				if (processor.changeState(this, COMPLETED)) {
+					processor.publisherDelegate.publishComplete();
+				}
 			}
-
+		},
+		RECEIVED {
 			@Override
 			public void writeComplete(AbstractResponseBodyFlushProcessor processor) {
+				try {
+					processor.flush();
+				}
+				catch (IOException ex) {
+					processor.cancel();
+					processor.onError(ex);
+				}
+
 				if (processor.subscriberCompleted) {
 					if (processor.changeState(this, COMPLETED)) {
 						processor.publisherDelegate.publishComplete();
 					}
 				}
 				else {
-					try {
-						processor.flush();
+					if (processor.changeState(this, REQUESTED)) {
 						processor.subscription.request(1);
 					}
-					catch (IOException ex) {
-						processor.cancel();
-						processor.onError(ex);
-					}
 				}
 			}
-		}, COMPLETED {
+
+			@Override
+			void onComplete(AbstractResponseBodyFlushProcessor processor) {
+				processor.subscriberCompleted = true;
+			}
+		},
+		COMPLETED {
 			@Override
 			public void onNext(AbstractResponseBodyFlushProcessor processor,
 					Publisher<DataBuffer> publisher) {
@@ -186,11 +201,6 @@ void onError(AbstractResponseBodyFlushProcessor processor, Throwable t) {
 			void onComplete(AbstractResponseBodyFlushProcessor processor) {
 				// ignore
 			}
-
-			@Override
-			public void writeComplete(AbstractResponseBodyFlushProcessor processor) {
-				// ignore
-			}
 		};
 
 		public void onSubscribe(AbstractResponseBodyFlushProcessor processor,
@@ -214,7 +224,7 @@ void onComplete(AbstractResponseBodyFlushProcessor processor) {
 		}
 
 		public void writeComplete(AbstractResponseBodyFlushProcessor processor) {
-			throw new IllegalStateException(toString());
+			// ignore
 		}
 
 		private static class WriteSubscriber implements Subscriber<Void> {

spring-web/src/test/java/org/springframework/http/server/reactive/WriteOnlyHandlerIntegrationTests.java

@@ -0,0 +1,86 @@
+/*
+ * Copyright 2002-2016 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.http.server.reactive;
+
+import java.net.URI;
+import java.nio.charset.StandardCharsets;
+import java.util.Random;
+
+import org.junit.Test;
+import org.springframework.core.io.buffer.DataBuffer;
+import org.springframework.http.RequestEntity;
+import org.springframework.http.ResponseEntity;
+import org.springframework.http.server.reactive.HttpHandler;
+import org.springframework.http.server.reactive.ServerHttpRequest;
+import org.springframework.http.server.reactive.ServerHttpResponse;
+import org.springframework.web.client.RestTemplate;
+
+import static org.junit.Assert.*;
+
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
+
+/**
+ * @author Violeta Georgieva
+ * @since 5.0
+ */
+public class WriteOnlyHandlerIntegrationTests extends AbstractHttpHandlerIntegrationTests {
+
+	private static final int REQUEST_SIZE = 4096 * 3;
+
+	private Random rnd = new Random();
+
+	private byte[] body;
+
+
+	@Override
+	protected WriteOnlyHandler createHttpHandler() {
+		return new WriteOnlyHandler();
+	}
+
+
+	@Test
+	public void writeOnly() throws Exception {
+		RestTemplate restTemplate = new RestTemplate();
+
+		this.body = randomBytes();
+		RequestEntity<byte[]> request = RequestEntity.post(
+				new URI("http://localhost:" + port)).body(
+						"".getBytes(StandardCharsets.UTF_8));
+		ResponseEntity<byte[]> response = restTemplate.exchange(request, byte[].class);
+
+		assertArrayEquals(body, response.getBody());
+	}
+
+
+	private byte[] randomBytes() {
+		byte[] buffer = new byte[REQUEST_SIZE];
+		rnd.nextBytes(buffer);
+		return buffer;
+	}
+
+
+	public class WriteOnlyHandler implements HttpHandler {
+
+		@Override
+		public Mono<Void> handle(ServerHttpRequest request, ServerHttpResponse response) {
+			DataBuffer buffer = response.bufferFactory().allocateBuffer(body.length);
+			buffer.write(body);
+			return response.writeAndFlushWith(Flux.just(Flux.just(buffer)));
+		}
+	}
+}