DATAMONGO-2414 - Guard drain loop in AsyncInputStreamHandler with state switch.

We now use a non-blocking state switch to determine whether to invoke drainLoop(…) from Subscriber completion.

Previously, we relied on same thread identification assuming if the subscription thread and the completion thread were the same, that we're already running inside the drain loop.
It turns out that a I/O thread could also run in event-loop mode where subscription and completion happens on the same thread but in between there's some processing and so the the call to completion is a delayed signal and not being called on the same stack as drainLoop(…).
The same-thread assumption was in place to avoid StackOverflow caused by infinite recursions.

We now use a state lock to enter the drain loop. Any concurrent attempts to re-enter the drain loop in Subscriber completion is now prevented to make sure that we continue draining while not causing stack recursions.

Original Pull Request: #807

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/gridfs/DataBufferPublisherAdapter.java

@@ -71,7 +71,6 @@ class DataBufferPublisherAdapter {
 
 		AsyncInputStreamHandler streamHandler = new AsyncInputStreamHandler(inputStream, inputStream.dataBufferFactory,
 				inputStream.bufferSize);
-
 		return Flux.create((sink) -> {
 
 			sink.onDispose(streamHandler::close);
@@ -146,12 +145,18 @@ class DataBufferPublisherAdapter {
 		private static final AtomicIntegerFieldUpdater<AsyncInputStreamHandler> STATE = AtomicIntegerFieldUpdater
 				.newUpdater(AsyncInputStreamHandler.class, "state");
 
+		private static final AtomicIntegerFieldUpdater<AsyncInputStreamHandler> DRAIN = AtomicIntegerFieldUpdater
+				.newUpdater(AsyncInputStreamHandler.class, "drain");
+
 		private static final AtomicIntegerFieldUpdater<AsyncInputStreamHandler> READ = AtomicIntegerFieldUpdater
 				.newUpdater(AsyncInputStreamHandler.class, "read");
 
 		private static final int STATE_OPEN = 0;
 		private static final int STATE_CLOSED = 1;
 
+		private static final int DRAIN_NONE = 0;
+		private static final int DRAIN_COMPLETION = 1;
+
 		private static final int READ_NONE = 0;
 		private static final int READ_IN_PROGRESS = 1;
 
@@ -165,6 +170,9 @@ class DataBufferPublisherAdapter {
 		// see STATE
 		volatile int state = STATE_OPEN;
 
+		// see DRAIN
+		volatile int drain = DRAIN_NONE;
+
 		// see READ_IN_PROGRESS
 		volatile int read = READ_NONE;
 
@@ -209,6 +217,14 @@ class DataBufferPublisherAdapter {
 			STATE.compareAndSet(this, STATE_OPEN, STATE_CLOSED);
 		}
 
+		boolean enterDrainLoop() {
+			return DRAIN.compareAndSet(this, DRAIN_NONE, DRAIN_COMPLETION);
+		}
+
+		void leaveDrainLoop() {
+			DRAIN.set(this, DRAIN_NONE);
+		}
+
 		boolean isClosed() {
 			return STATE.get(this) == STATE_CLOSED;
 		}
@@ -235,7 +251,6 @@ class DataBufferPublisherAdapter {
 			private final FluxSink<DataBuffer> sink;
 			private final DataBufferFactory factory;
 			private final ByteBuffer transport;
-			private final Thread subscribeThread = Thread.currentThread();
 			private volatile Subscription subscription;
 
 			BufferCoreSubscriber(FluxSink<DataBuffer> sink, DataBufferFactory factory, ByteBuffer transport) {
@@ -261,8 +276,6 @@ class DataBufferPublisherAdapter {
 			public void onNext(Integer bytes) {
 
 				if (isClosed()) {
-
-					onReadDone();
 					return;
 				}
 
@@ -273,14 +286,10 @@ class DataBufferPublisherAdapter {
 					decrementDemand();
 				}
 
-				try {
 				if (bytes == -1) {
 					sink.complete();
 					return;
 				}
-				} finally {
-					onReadDone();
-				}
 
 				subscription.request(1);
 			}
@@ -306,15 +315,25 @@ class DataBufferPublisherAdapter {
 					return;
 				}
 
-				onReadDone();
+				close();
 				sink.error(t);
 			}
 
 			@Override
 			public void onComplete() {
 
-				if (subscribeThread != Thread.currentThread()) {
+				onReadDone();
+
+				if (!isClosed()) {
+
+					if (enterDrainLoop()) {
+						try {
 							drainLoop(sink);
+						} finally {
+							leaveDrainLoop();
+						}
+					}
+
 				}
 			}
 		}