DATAMONGO-2393 - Use drain loop for same-thread processing in GridFS download stream.

We now rely on an outer drain-loop when GridFS reads complete on the same thread instead of using recursive subscriptions to avoid StackOverflow. Previously, we recursively invoked subscriptions that lead to an increased stack size.

Original Pull Request: #799

Author: mp911de

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

@@ -32,7 +32,6 @@
 
 import org.springframework.core.io.buffer.DataBuffer;
 import org.springframework.core.io.buffer.DataBufferFactory;
-import org.springframework.core.io.buffer.DataBufferUtils;
 
 import com.mongodb.reactivestreams.client.gridfs.AsyncInputStream;
 
@@ -51,11 +50,13 @@ class DataBufferPublisherAdapter {
 	 *
 	 * @param inputStream must not be {@literal null}.
 	 * @param dataBufferFactory must not be {@literal null}.
+	 * @param bufferSize read {@code n} bytes per iteration.
 	 * @return the resulting {@link Publisher}.
 	 */
-	static Flux<DataBuffer> createBinaryStream(AsyncInputStream inputStream, DataBufferFactory dataBufferFactory) {
+	static Flux<DataBuffer> createBinaryStream(AsyncInputStream inputStream, DataBufferFactory dataBufferFactory,
+			int bufferSize) {
 
-		State state = new State(inputStream, dataBufferFactory);
+		State state = new State(inputStream, dataBufferFactory, bufferSize);
 
 		return Flux.usingWhen(Mono.just(inputStream), it -> {
 
@@ -92,6 +93,7 @@ static class State {
 
 		final AsyncInputStream inputStream;
 		final DataBufferFactory dataBufferFactory;
+		final int bufferSize;
 
 		// see DEMAND
 		volatile long demand;
@@ -105,8 +107,16 @@ static class State {
 		void request(FluxSink<DataBuffer> sink, long n) {
 
 			Operators.addCap(DEMAND, this, n);
+			drainLoop(sink);
+		}
 
-			if (onShouldRead()) {
+		/**
+		 * Loops while we have demand and while no read is in progress.
+		 *
+		 * @param sink
+		 */
+		void drainLoop(FluxSink<DataBuffer> sink) {
+			while (onShouldRead()) {
 				emitNext(sink);
 			}
 		}
@@ -119,16 +129,16 @@ boolean onWantRead() {
 			return READ.compareAndSet(this, READ_NONE, READ_IN_PROGRESS);
 		}
 
-		boolean onReadDone() {
-			return READ.compareAndSet(this, READ_IN_PROGRESS, READ_NONE);
+		void onReadDone() {
+			READ.compareAndSet(this, READ_IN_PROGRESS, READ_NONE);
 		}
 
 		long getDemand() {
 			return DEMAND.get(this);
 		}
 
-		boolean decrementDemand() {
-			return DEMAND.decrementAndGet(this) > 0;
+		void decrementDemand() {
+			DEMAND.decrementAndGet(this);
 		}
 
 		void close() {
@@ -143,30 +153,32 @@ boolean isClosed() {
 		 * Emit the next {@link DataBuffer}.
 		 *
 		 * @param sink
+		 * @return
 		 */
-		void emitNext(FluxSink<DataBuffer> sink) {
-
-			DataBuffer dataBuffer = dataBufferFactory.allocateBuffer();
-			ByteBuffer intermediate = ByteBuffer.allocate(dataBuffer.capacity());
+		private void emitNext(FluxSink<DataBuffer> sink) {
 
+			ByteBuffer transport = ByteBuffer.allocate(bufferSize);
+			BufferCoreSubscriber bufferCoreSubscriber = new BufferCoreSubscriber(sink, dataBufferFactory, transport);
 			try {
-				Mono.from(inputStream.read(intermediate)).subscribe(new BufferCoreSubscriber(sink, dataBuffer, intermediate));
-			} catch (Exception e) {
+				inputStream.read(transport).subscribe(bufferCoreSubscriber);
+			} catch (Throwable e) {
 				sink.error(e);
 			}
 		}
 
 		private class BufferCoreSubscriber implements CoreSubscriber<Integer> {
 
 			private final FluxSink<DataBuffer> sink;
-			private final DataBuffer dataBuffer;
-			private final ByteBuffer intermediate;
+			private final DataBufferFactory factory;
+			private final ByteBuffer transport;
+			private final Thread subscribeThread = Thread.currentThread();
+			private volatile Subscription subscription;
 
-			BufferCoreSubscriber(FluxSink<DataBuffer> sink, DataBuffer dataBuffer, ByteBuffer intermediate) {
+			BufferCoreSubscriber(FluxSink<DataBuffer> sink, DataBufferFactory factory, ByteBuffer transport) {
 
 				this.sink = sink;
-				this.dataBuffer = dataBuffer;
-				this.intermediate = intermediate;
+				this.factory = factory;
+				this.transport = transport;
 			}
 
 			@Override
@@ -176,6 +188,7 @@ public Context currentContext() {
 
 			@Override
 			public void onSubscribe(Subscription s) {
+				this.subscription = s;
 				s.request(1);
 			}
 
@@ -185,24 +198,32 @@ public void onNext(Integer bytes) {
 				if (isClosed()) {
 
 					onReadDone();
-					DataBufferUtils.release(dataBuffer);
-					Operators.onNextDropped(dataBuffer, sink.currentContext());
 					return;
 				}
 
-				intermediate.flip();
-				dataBuffer.write(intermediate);
+				if (bytes > 0) {
 
-				sink.next(dataBuffer);
-				decrementDemand();
+					transport.flip();
+
+					DataBuffer dataBuffer = factory.allocateBuffer(transport.remaining());
+					dataBuffer.write(transport);
+
+					transport.clear();
+					sink.next(dataBuffer);
+
+					decrementDemand();
+				}
 
 				try {
 					if (bytes == -1) {
 						sink.complete();
+						return;
 					}
 				} finally {
 					onReadDone();
 				}
+
+				subscription.request(1);
 			}
 
 			@Override
@@ -215,16 +236,14 @@ public void onError(Throwable t) {
 				}
 
 				onReadDone();
-				DataBufferUtils.release(dataBuffer);
-				Operators.onNextDropped(dataBuffer, sink.currentContext());
 				sink.error(t);
 			}
 
 			@Override
 			public void onComplete() {
 
-				if (onShouldRead()) {
-					emitNext(sink);
+				if (subscribeThread != Thread.currentThread()) {
+					drainLoop(sink);
 				}
 			}
 		}