Merge pull request #1123 from violetagg/violetagg-flush-last-element

Author: rstoyanchev

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

@@ -26,59 +26,47 @@
 import org.springframework.core.io.buffer.DataBufferFactory;
 
 /**
- * Abstract base class for listener-based server responses, i.e. Servlet 3.1 and Undertow.
+ * Abstract base class for listener-based server responses, e.g. Servlet 3.1
+ * and Undertow.
+ *
  * @author Arjen Poutsma
  * @since 5.0
  */
 public abstract class AbstractListenerServerHttpResponse extends AbstractServerHttpResponse {
 
 	private final AtomicBoolean writeCalled = new AtomicBoolean();
 
+
 	public AbstractListenerServerHttpResponse(DataBufferFactory dataBufferFactory) {
 		super(dataBufferFactory);
 	}
 
+
 	@Override
 	protected final Mono<Void> writeWithInternal(Publisher<DataBuffer> body) {
-		if (this.writeCalled.compareAndSet(false, true)) {
-			Processor<DataBuffer, Void> bodyProcessor = createBodyProcessor();
-			return Mono.from(subscriber -> {
-				body.subscribe(bodyProcessor);
-				bodyProcessor.subscribe(subscriber);
-			});
-
-		} else {
-			return Mono.error(new IllegalStateException(
-					"writeWith() or writeAndFlushWith() has already been called"));
-		}
+		return writeAndFlushWithInternal(Mono.just(body));
 	}
 
 	@Override
 	protected final Mono<Void> writeAndFlushWithInternal(Publisher<Publisher<DataBuffer>> body) {
 		if (this.writeCalled.compareAndSet(false, true)) {
-			Processor<Publisher<DataBuffer>, Void> bodyProcessor =
-					createBodyFlushProcessor();
+			Processor<Publisher<DataBuffer>, Void> bodyProcessor = createBodyFlushProcessor();
 			return Mono.from(subscriber -> {
 				body.subscribe(bodyProcessor);
 				bodyProcessor.subscribe(subscriber);
 			});
-		} else {
+		}
+		else {
 			return Mono.error(new IllegalStateException(
 					"writeWith() or writeAndFlushWith() has already been called"));
 		}
 	}
 
-	/**
-	 * Abstract template method to create a {@code Processor<DataBuffer, Void>} that
-	 * will write the response body to the underlying output. Called from
-	 * {@link #writeWithInternal(Publisher)}.
-	 */
-	protected abstract Processor<DataBuffer, Void> createBodyProcessor();
-
 	/**
 	 * Abstract template method to create a {@code Processor<Publisher<DataBuffer>, Void>}
 	 * that will write the response body with flushes to the underlying output. Called from
 	 * {@link #writeAndFlushWithInternal(Publisher)}.
 	 */
 	protected abstract Processor<Publisher<DataBuffer>, Void> createBodyFlushProcessor();
+
 }

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

@@ -35,26 +35,25 @@
  * Servlet 3.1 and Undertow support.
  *
  * @author Arjen Poutsma
+ * @author Violeta Georgieva
  * @since 5.0
  * @see ServletServerHttpRequest
  * @see UndertowHttpHandlerAdapter
  * @see ServerHttpResponse#writeAndFlushWith(Publisher)
  */
-abstract class AbstractResponseBodyFlushProcessor
-		implements Processor<Publisher<DataBuffer>, Void> {
+abstract class AbstractResponseBodyFlushProcessor implements Processor<Publisher<DataBuffer>, Void> {
 
 	protected final Log logger = LogFactory.getLog(getClass());
 
-	private final ResponseBodyWriteResultPublisher publisherDelegate =
-			new ResponseBodyWriteResultPublisher();
+	private final ResponseBodyWriteResultPublisher resultPublisher = new ResponseBodyWriteResultPublisher();
 
-	private final AtomicReference<State> state =
-			new AtomicReference<>(State.UNSUBSCRIBED);
+	private final AtomicReference<State> state = new AtomicReference<>(State.UNSUBSCRIBED);
 
 	private volatile boolean subscriberCompleted;
 
 	private Subscription subscription;
 
+
 	// Subscriber
 
 	@Override
@@ -89,13 +88,15 @@ public final void onComplete() {
 		this.state.get().onComplete(this);
 	}
 
+
 	// Publisher
 
 	@Override
 	public final void subscribe(Subscriber<? super Void> subscriber) {
-		this.publisherDelegate.subscribe(subscriber);
+		this.resultPublisher.subscribe(subscriber);
 	}
 
+
 	/**
 	 * Creates a new processor for subscribing to a body chunk.
 	 */
@@ -106,6 +107,11 @@ public final void subscribe(Subscriber<? super Void> subscriber) {
 	 */
 	protected abstract void flush() throws IOException;
 
+
+	private boolean changeState(State oldState, State newState) {
+		return this.state.compareAndSet(oldState, newState);
+	}
+
 	private void writeComplete() {
 		if (logger.isTraceEnabled()) {
 			logger.trace(this.state + " writeComplete");
@@ -114,104 +120,117 @@ private void writeComplete() {
 
 	}
 
-	private boolean changeState(State oldState, State newState) {
-		return this.state.compareAndSet(oldState, newState);
+	private void cancel() {
+		this.subscription.cancel();
 	}
 
+
 	private enum State {
+
 		UNSUBSCRIBED {
+
 			@Override
-			public void onSubscribe(AbstractResponseBodyFlushProcessor processor,
-					Subscription subscription) {
+			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));
+			public void onNext(AbstractResponseBodyFlushProcessor processor, Publisher<DataBuffer> chunk) {
+				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;
+			public void onComplete(AbstractResponseBodyFlushProcessor processor) {
+				if (processor.changeState(this, COMPLETED)) {
+					processor.resultPublisher.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();
+						processor.resultPublisher.publishComplete();
 					}
 				}
 				else {
-					try {
-						processor.flush();
-					}
-					catch (IOException ex) {
-						processor.onError(ex);
+					if (processor.changeState(this, REQUESTED)) {
+						processor.subscription.request(1);
 					}
-					processor.subscription.request(1);
 				}
 			}
-		}, COMPLETED {
-			@Override
-			public void onNext(AbstractResponseBodyFlushProcessor processor,
-					Publisher<DataBuffer> publisher) {
-				// ignore
 
+			@Override
+			public void onComplete(AbstractResponseBodyFlushProcessor processor) {
+				processor.subscriberCompleted = true;
 			}
+		},
+		COMPLETED {
 
 			@Override
-			void onError(AbstractResponseBodyFlushProcessor processor, Throwable t) {
+			public void onNext(AbstractResponseBodyFlushProcessor processor,
+					Publisher<DataBuffer> publisher) {
 				// ignore
+
 			}
 
 			@Override
-			void onComplete(AbstractResponseBodyFlushProcessor processor) {
+			public void onError(AbstractResponseBodyFlushProcessor processor, Throwable t) {
 				// ignore
 			}
 
 			@Override
-			public void writeComplete(AbstractResponseBodyFlushProcessor processor) {
+			public void onComplete(AbstractResponseBodyFlushProcessor processor) {
 				// ignore
 			}
 		};
 
-		public void onSubscribe(AbstractResponseBodyFlushProcessor processor,
-				Subscription subscription) {
+		public void onSubscribe(AbstractResponseBodyFlushProcessor processor, Subscription subscription) {
 			subscription.cancel();
 		}
 
-		public void onNext(AbstractResponseBodyFlushProcessor processor,
-				Publisher<DataBuffer> publisher) {
+		public void onNext(AbstractResponseBodyFlushProcessor processor, Publisher<DataBuffer> publisher) {
 			throw new IllegalStateException(toString());
 		}
 
-		void onError(AbstractResponseBodyFlushProcessor processor, Throwable t) {
+		public void onError(AbstractResponseBodyFlushProcessor processor, Throwable ex) {
 			if (processor.changeState(this, COMPLETED)) {
-				processor.publisherDelegate.publishError(t);
+				processor.resultPublisher.publishError(ex);
 			}
 		}
 
-		void onComplete(AbstractResponseBodyFlushProcessor processor) {
+		public void onComplete(AbstractResponseBodyFlushProcessor processor) {
 			throw new IllegalStateException(toString());
 		}
 
 		public void writeComplete(AbstractResponseBodyFlushProcessor processor) {
-			throw new IllegalStateException(toString());
+			// ignore
 		}
 
+
 		private static class WriteSubscriber implements Subscriber<Void> {
 
 			private final AbstractResponseBodyFlushProcessor processor;
@@ -221,22 +240,23 @@ public WriteSubscriber(AbstractResponseBodyFlushProcessor processor) {
 			}
 
 			@Override
-			public void onSubscribe(Subscription s) {
-				s.request(Long.MAX_VALUE);
+			public void onSubscribe(Subscription subscription) {
+				subscription.request(Long.MAX_VALUE);
 			}
 
 			@Override
 			public void onNext(Void aVoid) {
 			}
 
 			@Override
-			public void onError(Throwable t) {
-				processor.onError(t);
+			public void onError(Throwable ex) {
+				this.processor.cancel();
+				this.processor.onError(ex);
 			}
 
 			@Override
 			public void onComplete() {
-				processor.writeComplete();
+				this.processor.writeComplete();
 			}
 		}
 	}