Set readyToWrite flag after cached signals emitted

Issue: SPR-16555

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2017 the original author or authors.
+ * Copyright 2002-2018 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.
@@ -75,6 +75,33 @@ public void subscribe(CoreSubscriber<? super Void> actual) {
 	}
 
 
+	private enum State {
+
+		/** No emissions from the upstream source yet */
+		NEW,
+
+		/**
+		 * At least one signal of any kind has been received; we're ready to
+		 * call the write function and proceed with actual writing.
+		 */
+		FIRST_SIGNAL_RECEIVED,
+
+		/**
+		 * The write subscriber has subscribed and requested; we're going to
+		 * emit the cached signals.
+		 */
+		EMITTING_CACHED_SIGNALS,
+
+		/**
+		 * The write subscriber has subscribed, and cached signals have been
+		 * emitted to it; we're ready to switch to a simple pass-through mode
+		 * for all remaining signals.
+		 **/
+		READY_TO_WRITE
+
+	}
+
+
 	/**
 	 * A barrier inserted between the write source and the write subscriber
 	 * (i.e. the HTTP server adapter) that pre-fetches and waits for the first
@@ -99,27 +126,23 @@ private class WriteBarrier implements CoreSubscriber<T>, Subscription, Publisher
 		@Nullable
 		private Subscription subscription;
 
-		/**
-		 * We've at at least one emission, we've called the write function, the write
-		 * subscriber has subscribed and cached signals have been emitted to it.
-		 * We're now simply passing data through to the write subscriber.
-		 **/
-		private boolean readyToWrite = false;
-
-		/** No emission from upstream yet */
-		private boolean beforeFirstEmission = true;
-
-		/** Cached signal before readyToWrite */
+		/** Cached data item before readyToWrite */
 		@Nullable
 		private T item;
 
-		/** Cached 1st/2nd signal before readyToWrite */
+		/** Cached error signal before readyToWrite */
 		@Nullable
 		private Throwable error;
 
-		/** Cached 1st/2nd signal before readyToWrite */
+		/** Cached onComplete signal before readyToWrite */
 		private boolean completed = false;
 
+		/** Recursive demand while emitting cached signals */
+		private long demandBeforeReadyToWrite;
+
+		/** Current state */
+		private State state = State.NEW;
+
 		/** The actual writeSubscriber from the HTTP server adapter */
 		@Nullable
 		private Subscriber<? super T> writeSubscriber;
@@ -143,18 +166,18 @@ public final void onSubscribe(Subscription s) {
 
 		@Override
 		public final void onNext(T item) {
-			if (this.readyToWrite) {
+			if (this.state == State.READY_TO_WRITE) {
 				requiredWriteSubscriber().onNext(item);
 				return;
 			}
 			//FIXME revisit in case of reentrant sync deadlock
 			synchronized (this) {
-				if (this.readyToWrite) {
+				if (this.state == State.READY_TO_WRITE) {
 					requiredWriteSubscriber().onNext(item);
 				}
-				else if (this.beforeFirstEmission) {
+				else if (this.state == State.NEW) {
 					this.item = item;
-					this.beforeFirstEmission = false;
+					this.state = State.FIRST_SIGNAL_RECEIVED;
 					writeFunction.apply(this).subscribe(this.writeCompletionBarrier);
 				}
 				else {
@@ -173,16 +196,16 @@ private Subscriber<? super T> requiredWriteSubscriber() {
 
 		@Override
 		public final void onError(Throwable ex) {
-			if (this.readyToWrite) {
+			if (this.state == State.READY_TO_WRITE) {
 				requiredWriteSubscriber().onError(ex);
 				return;
 			}
 			synchronized (this) {
-				if (this.readyToWrite) {
+				if (this.state == State.READY_TO_WRITE) {
 					requiredWriteSubscriber().onError(ex);
 				}
-				else if (this.beforeFirstEmission) {
-					this.beforeFirstEmission = false;
+				else if (this.state == State.NEW) {
+					this.state = State.FIRST_SIGNAL_RECEIVED;
 					this.writeCompletionBarrier.onError(ex);
 				}
 				else {
@@ -193,17 +216,17 @@ else if (this.beforeFirstEmission) {
 
 		@Override
 		public final void onComplete() {
-			if (this.readyToWrite) {
+			if (this.state == State.READY_TO_WRITE) {
 				requiredWriteSubscriber().onComplete();
 				return;
 			}
 			synchronized (this) {
-				if (this.readyToWrite) {
+				if (this.state == State.READY_TO_WRITE) {
 					requiredWriteSubscriber().onComplete();
 				}
-				else if (this.beforeFirstEmission) {
+				else if (this.state == State.NEW) {
 					this.completed = true;
-					this.beforeFirstEmission = false;
+					this.state = State.FIRST_SIGNAL_RECEIVED;
 					writeFunction.apply(this).subscribe(this.writeCompletionBarrier);
 				}
 				else {
@@ -226,19 +249,28 @@ public void request(long n) {
 			if (s == null) {
 				return;
 			}
-			if (this.readyToWrite) {
+			if (this.state == State.READY_TO_WRITE) {
 				s.request(n);
 				return;
 			}
 			synchronized (this) {
 				if (this.writeSubscriber != null) {
-					this.readyToWrite = true;
-					if (emitCachedSignals()) {
+					if (this.state == State.EMITTING_CACHED_SIGNALS) {
+						this.demandBeforeReadyToWrite = n;
 						return;
 					}
-					n--;
-					if (n == 0) {
-						return;
+					try {
+						this.state = State.EMITTING_CACHED_SIGNALS;
+						if (emitCachedSignals()) {
+							return;
+						}
+						n = n + this.demandBeforeReadyToWrite - 1;
+						if (n == 0) {
+							return;
+						}
+					}
+					finally {
+						this.state = State.READY_TO_WRITE;
 					}
 				}
 			}