Support Publisher to InputStream conversion

See gh-31677

spring-core/src/main/java/org/springframework/core/io/buffer/DataBufferUtils.java

@@ -456,6 +456,35 @@ public abstract class DataBufferUtils {
 				consumer::accept, new DataBufferMapper(bufferFactory), executor, chunkSize);
 	}
 
+	/**
+	 * Subscribes to given {@link Publisher} and returns subscription
+	 * as {@link InputStream} that allows reading all propagated {@link DataBuffer} messages via its imperative API.
+	 * Given the {@link InputStream} implementation buffers messages as per configuration.
+	 * The returned {@link InputStream} is considered terminated when the given {@link Publisher} signaled one of the
+	 * terminal signal ({@link Subscriber#onComplete() or {@link Subscriber#onError(Throwable)}})
+	 * and all the stored {@link DataBuffer} polled from the internal buffer.
+	 * The returned {@link InputStream} will call {@link Subscription#cancel()} and release all stored {@link DataBuffer}
+	 * when {@link InputStream#close()} is called.
+	 * <p>
+	 * Note: The implementation of the returned {@link InputStream} disallow concurrent call on
+	 * any of the {@link InputStream#read} methods
+	 * <p>
+	 * Note: {@link Subscription#request(long)} happens eagerly for the first time upon subscription
+	 * and then repeats every time {@code bufferSize - (bufferSize >> 2)} consumed
+	 *
+	 * @param publisher the source of {@link DataBuffer} which should be represented as an {@link InputStream}
+	 * @param bufferSize the maximum amount of {@link DataBuffer} prefetched in advance and stored inside {@link InputStream}
+	 * @return an {@link InputStream} instance representing given {@link Publisher} messages
+	 */
+	public static <T extends DataBuffer> InputStream subscribeAsInputStream(Publisher<T> publisher, int bufferSize) {
+		Assert.notNull(publisher, "Publisher must not be null");
+		Assert.isTrue(bufferSize > 0, "Buffer size must be > 0");
+
+		InputStreamSubscriber inputStreamSubscriber = new InputStreamSubscriber(bufferSize);
+		publisher.subscribe(inputStreamSubscriber);
+		return inputStreamSubscriber;
+	}
+
 
 	//---------------------------------------------------------------------
 	// Various

spring-core/src/main/java/org/springframework/core/io/buffer/InputStreamSubscriber.java

@@ -0,0 +1,355 @@
+package org.springframework.core.io.buffer;
+
+import org.reactivestreams.Publisher;
+import org.reactivestreams.Subscriber;
+import org.reactivestreams.Subscription;
+import org.springframework.lang.Nullable;
+import reactor.core.Exceptions;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.ByteBuffer;
+import java.util.ConcurrentModificationException;
+import java.util.Objects;
+import java.util.Queue;
+import java.util.concurrent.ArrayBlockingQueue;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.concurrent.locks.LockSupport;
+import java.util.concurrent.locks.ReentrantLock;
+
+/**
+ * Bridges between {@link Publisher Publisher<DataBuffer>} and {@link InputStream}.
+ *
+ * <p>Note that this class has a near duplicate in
+ * {@link org.springframework.http.client.InputStreamSubscriber}.
+ *
+ * @author Oleh Dokuka
+ * @since 6.1
+ */
+final class InputStreamSubscriber extends InputStream implements Subscriber<DataBuffer> {
+
+	static final Object READY = new Object();
+	static final DataBuffer DONE = DefaultDataBuffer.fromEmptyByteBuffer(DefaultDataBufferFactory.sharedInstance, ByteBuffer.allocate(0));
+	static final DataBuffer CLOSED = DefaultDataBuffer.fromEmptyByteBuffer(DefaultDataBufferFactory.sharedInstance, ByteBuffer.allocate(0));
+
+	final int prefetch;
+	final int limit;
+	final ReentrantLock       lock;
+	final Queue<DataBuffer> queue;
+
+	final AtomicReference<Object> parkedThread = new AtomicReference<>();
+	final AtomicInteger workAmount = new AtomicInteger();
+
+	volatile boolean closed;
+	int consumed;
+
+	@Nullable
+	DataBuffer available;
+
+	@Nullable
+	Subscription s;
+	boolean done;
+	@Nullable
+	Throwable error;
+
+	InputStreamSubscriber(int prefetch) {
+		this.prefetch = prefetch;
+		this.limit = prefetch == Integer.MAX_VALUE ? Integer.MAX_VALUE : prefetch - (prefetch >> 2);
+		this.queue = new ArrayBlockingQueue<>(prefetch);
+		this.lock = new ReentrantLock(false);
+	}
+
+	@Override
+	public void onSubscribe(Subscription subscription) {
+		if (this.s != null) {
+			subscription.cancel();
+			return;
+		}
+
+		this.s = subscription;
+		subscription.request(prefetch == Integer.MAX_VALUE ? Long.MAX_VALUE : prefetch);
+	}
+
+	@Override
+	public void onNext(DataBuffer t) {
+		if (this.done) {
+			discard(t);
+			return;
+		}
+
+		if (!queue.offer(t)) {
+			discard(t);
+			error = new RuntimeException("Buffer overflow");
+			done = true;
+		}
+
+		int previousWorkState = addWork();
+		if (previousWorkState == Integer.MIN_VALUE) {
+			DataBuffer value = queue.poll();
+			if (value != null) {
+				discard(value);
+			}
+			return;
+		}
+
+		if (previousWorkState == 0) {
+			resume();
+		}
+	}
+
+	@Override
+	public void onError(Throwable throwable) {
+		if (this.done) {
+			return;
+		}
+		this.error = throwable;
+		this.done = true;
+
+		if (addWork() == 0) {
+			resume();
+		}
+	}
+
+	@Override
+	public void onComplete() {
+		if (this.done) {
+			return;
+		}
+
+		this.done = true;
+
+		if (addWork() == 0) {
+			resume();
+		}
+	}
+
+	int addWork() {
+		for (;;) {
+			int produced = this.workAmount.getPlain();
+
+			if (produced == Integer.MIN_VALUE) {
+				return Integer.MIN_VALUE;
+			}
+
+			int nextProduced = produced == Integer.MAX_VALUE ? 1 : produced + 1;
+
+
+			if (workAmount.weakCompareAndSetRelease(produced, nextProduced)) {
+				return produced;
+			}
+		}
+	}
+
+	@Override
+	public int read() throws IOException {
+		if (!lock.tryLock()) {
+			if (this.closed) {
+				return -1;
+			}
+			throw new ConcurrentModificationException("concurrent access is disallowed");
+		}
+
+		try {
+			DataBuffer bytes = getBytesOrAwait();
+
+			if (bytes == DONE) {
+				this.closed = true;
+				cleanAndFinalize();
+				if (this.error == null) {
+					return -1;
+				}
+				else {
+					throw Exceptions.propagate(error);
+				}
+			} else if (bytes == CLOSED) {
+				cleanAndFinalize();
+				return -1;
+			}
+
+			return bytes.read() & 0xFF;
+		}
+		catch (Throwable t) {
+			this.closed = true;
+			this.s.cancel();
+			cleanAndFinalize();
+			throw Exceptions.propagate(t);
+		}
+		finally {
+			lock.unlock();
+		}
+	}
+
+	@Override
+	public int read(byte[] b, int off, int len) throws IOException {
+		Objects.checkFromIndexSize(off, len, b.length);
+		if (len == 0) {
+			return 0;
+		}
+
+		if (!lock.tryLock()) {
+			if (this.closed) {
+				return -1;
+			}
+			throw new ConcurrentModificationException("concurrent access is disallowed");
+		}
+
+		try {
+			for (int j = 0; j < len;) {
+				DataBuffer bytes = getBytesOrAwait();
+
+				if (bytes == DONE) {
+					cleanAndFinalize();
+					if (this.error == null) {
+						this.closed = true;
+						return j == 0 ? -1 : j;
+					}
+					else {
+						if (j == 0) {
+							this.closed = true;
+							throw Exceptions.propagate(error);
+						}
+
+						return j;
+					}
+				} else if (bytes == CLOSED) {
+					this.s.cancel();
+					cleanAndFinalize();
+					return -1;
+				}
+				int initialReadPosition = bytes.readPosition();
+				bytes.read(b, off + j, Math.min(len - j, bytes.readableByteCount()));
+				j += bytes.readPosition() - initialReadPosition;
+			}
+
+			return len;
+		}
+		catch (Throwable t) {
+			this.closed = true;
+			this.s.cancel();
+			cleanAndFinalize();
+			throw Exceptions.propagate(t);
+		}
+		finally {
+			lock.unlock();
+		}
+	}
+
+	DataBuffer getBytesOrAwait() {
+		if (this.available == null || this.available.readableByteCount() == 0) {
+
+			discard(this.available);
+			this.available = null;
+
+			int actualWorkAmount = this.workAmount.getAcquire();
+			for (;;) {
+				if (this.closed) {
+					return CLOSED;
+				}
+
+				boolean d = this.done;
+				DataBuffer t = this.queue.poll();
+				if (t != null) {
+					int consumed = ++this.consumed;
+					this.available = t;
+					if (consumed == this.limit) {
+						this.consumed = 0;
+						this.s.request(this.limit);
+					}
+					break;
+				}
+
+				if (d) {
+					return DONE;
+				}
+
+				actualWorkAmount = workAmount.addAndGet(-actualWorkAmount);
+				if (actualWorkAmount == 0) {
+					await();
+				}
+			}
+		}
+
+		return this.available;
+	}
+
+	void cleanAndFinalize() {
+		discard(this.available);
+		this.available = null;
+
+		for (;;) {
+			int workAmount = this.workAmount.getPlain();
+			DataBuffer value;
+
+			while((value = queue.poll()) != null) {
+				discard(value);
+			}
+
+			if (this.workAmount.weakCompareAndSetPlain(workAmount, Integer.MIN_VALUE)) {
+				return;
+			}
+		}
+	}
+
+	void discard(@Nullable DataBuffer value) {
+		DataBufferUtils.release(value);
+	}
+
+	@Override
+	public void close() throws IOException {
+		if (this.closed) {
+			return;
+		}
+
+		this.closed = true;
+
+		if (!this.lock.tryLock()) {
+			if (addWork() == 0) {
+				resume();
+			}
+			return;
+		}
+
+		try {
+			this.s.cancel();
+			cleanAndFinalize();
+		}
+		finally {
+			this.lock.unlock();
+		}
+	}
+
+	private void await() {
+		Thread toUnpark = Thread.currentThread();
+
+		while (true) {
+			Object current = this.parkedThread.get();
+			if (current == READY) {
+				break;
+			}
+
+			if (current != null && current != toUnpark) {
+				throw new IllegalStateException("Only one (Virtual)Thread can await!");
+			}
+
+			if (parkedThread.compareAndSet( null, toUnpark)) {
+				LockSupport.park();
+				// we don't just break here because park() can wake up spuriously
+				// if we got a proper resume, get() == READY and the loop will quit above
+			}
+		}
+		// clear the resume indicator so that the next await call will park without a resume()
+		this.parkedThread.lazySet(null);
+	}
+
+	private void resume() {
+		if (this.parkedThread != READY) {
+			Object old = parkedThread.getAndSet(READY);
+			if (old != READY) {
+				LockSupport.unpark((Thread)old);
+			}
+		}
+	}
+
+
+}

spring-core/src/test/java/org/springframework/core/io/buffer/DataBufferUtilsTests.java

@@ -17,6 +17,7 @@
 package org.springframework.core.io.buffer;
 
 import java.io.IOException;
+import java.io.InputStream;
 import java.io.OutputStream;
 import java.io.OutputStreamWriter;
 import java.net.URI;
@@ -27,15 +28,18 @@ import java.nio.channels.FileChannel;
 import java.nio.channels.ReadableByteChannel;
 import java.nio.channels.SeekableByteChannel;
 import java.nio.channels.WritableByteChannel;
+import java.nio.charset.Charset;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.nio.file.StandardOpenOption;
 import java.time.Duration;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.Executors;
+import java.util.concurrent.ThreadLocalRandom;
 
 import io.netty.buffer.ByteBuf;
 import io.netty.buffer.PooledByteBufAllocator;
@@ -688,6 +692,189 @@ class DataBufferUtilsTests extends AbstractDataBufferAllocatingTests {
 		latch.await();
 	}
 
+
+
+	@ParameterizedDataBufferAllocatingTest
+	void inputStreamSubscriberChunkSize(DataBufferFactory bufferFactory) {
+		genericInputStreamSubscriberTest(bufferFactory, 3, 3, 64, List.of("foo", "bar", "baz"), List.of("foo", "bar", "baz"));
+	}
+
+	@ParameterizedDataBufferAllocatingTest
+	void inputStreamSubscriberChunkSize2(DataBufferFactory bufferFactory) {
+		genericInputStreamSubscriberTest(bufferFactory, 3, 3, 1, List.of("foo", "bar", "baz"), List.of("foo", "bar", "baz"));
+	}
+
+	@ParameterizedDataBufferAllocatingTest
+	void inputStreamSubscriberChunkSize3(DataBufferFactory bufferFactory) {
+		genericInputStreamSubscriberTest(bufferFactory, 3, 12, 1, List.of("foo", "bar", "baz"), List.of("foobarbaz"));
+	}
+
+	@ParameterizedDataBufferAllocatingTest
+	void inputStreamSubscriberChunkSize4(DataBufferFactory bufferFactory) {
+		genericInputStreamSubscriberTest(bufferFactory, 3, 1, 1, List.of("foo", "bar", "baz"), List.of("f", "o", "o", "b", "a", "r", "b", "a", "z"));
+	}
+
+	@ParameterizedDataBufferAllocatingTest
+	void inputStreamSubscriberChunkSize5(DataBufferFactory bufferFactory) {
+		genericInputStreamSubscriberTest(bufferFactory, 3, 2, 1, List.of("foo", "bar", "baz"), List.of("fo", "ob", "ar", "ba", "z"));
+	}
+
+	@ParameterizedDataBufferAllocatingTest
+	void inputStreamSubscriberChunkSize6(DataBufferFactory bufferFactory) {
+		genericInputStreamSubscriberTest(bufferFactory, 1, 3, 1, List.of("foo", "bar", "baz"), List.of("foo", "bar", "baz"));
+	}
+
+	@ParameterizedDataBufferAllocatingTest
+	void inputStreamSubscriberChunkSize7(DataBufferFactory bufferFactory) {
+		genericInputStreamSubscriberTest(bufferFactory, 1, 3, 64, List.of("foo", "bar", "baz"), List.of("foo", "bar", "baz"));
+	}
+
+	void genericInputStreamSubscriberTest(DataBufferFactory bufferFactory, int writeChunkSize, int readChunkSize, int bufferSize, List<String> input, List<String> expectedOutput) {
+		super.bufferFactory = bufferFactory;
+
+		Publisher<DataBuffer> publisher = DataBufferUtils.outputStreamPublisher(outputStream -> {
+			try {
+				for (String word : input) {
+					outputStream.write(word.getBytes(StandardCharsets.UTF_8));
+				}
+			}
+			catch (IOException ex) {
+				fail(ex.getMessage(), ex);
+			}
+		}, super.bufferFactory, Executors.newSingleThreadExecutor(), writeChunkSize);
+
+
+
+		byte[] chunk = new byte[readChunkSize];
+		ArrayList<String> words = new ArrayList<>();
+
+		try (InputStream inputStream = DataBufferUtils.subscribeAsInputStream(publisher, bufferSize)) {
+			int read;
+			while((read = inputStream.read(chunk)) > -1) {
+				String word = new String(chunk, 0, read, StandardCharsets.UTF_8);
+				words.add(word);
+			}
+		}
+		catch (IOException e) {
+			throw new RuntimeException(e);
+		}
+		assertThat(words).containsExactlyElementsOf(expectedOutput);
+	}
+
+	@ParameterizedDataBufferAllocatingTest
+	void inputStreamSubscriberError(DataBufferFactory bufferFactory) throws InterruptedException {
+		super.bufferFactory = bufferFactory;
+
+		var input = List.of("foo ", "bar ", "baz");
+
+		Publisher<DataBuffer> publisher = DataBufferUtils.outputStreamPublisher(outputStream -> {
+			try {
+				for (String word : input) {
+					outputStream.write(word.getBytes(StandardCharsets.UTF_8));
+				}
+				throw new RuntimeException("boom");
+			}
+			catch (IOException ex) {
+				fail(ex.getMessage(), ex);
+			}
+		}, super.bufferFactory, Executors.newSingleThreadExecutor(), 1);
+
+
+		RuntimeException error = null;
+		byte[] chunk = new byte[4];
+		ArrayList<String> words = new ArrayList<>();
+
+		try (InputStream inputStream = DataBufferUtils.subscribeAsInputStream(publisher, 1)) {
+			int read;
+			while((read = inputStream.read(chunk)) > -1) {
+				String word = new String(chunk, 0, read, StandardCharsets.UTF_8);
+				words.add(word);
+			}
+		}
+		catch (IOException e) {
+			throw new RuntimeException(e);
+		}
+		catch (RuntimeException e) {
+			error = e;
+		}
+		assertThat(words).containsExactlyElementsOf(List.of("foo ", "bar ", "baz"));
+		assertThat(error).hasMessage("boom");
+	}
+
+	@ParameterizedDataBufferAllocatingTest
+	void inputStreamSubscriberMixedReadMode(DataBufferFactory bufferFactory) throws InterruptedException {
+		super.bufferFactory = bufferFactory;
+
+		var input = List.of("foo ", "bar ", "baz");
+
+		Publisher<DataBuffer> publisher = DataBufferUtils.outputStreamPublisher(outputStream -> {
+			try {
+				for (String word : input) {
+					outputStream.write(word.getBytes(StandardCharsets.UTF_8));
+				}
+			}
+			catch (IOException ex) {
+				fail(ex.getMessage(), ex);
+			}
+		}, super.bufferFactory, Executors.newSingleThreadExecutor(), 1);
+
+
+		byte[] chunk = new byte[3];
+		ArrayList<String> words = new ArrayList<>();
+
+		try (InputStream inputStream = DataBufferUtils.subscribeAsInputStream(publisher, 1)) {
+			words.add(new String(chunk,0, inputStream.read(chunk), StandardCharsets.UTF_8));
+			assertThat(inputStream.read()).isEqualTo(' ' & 0xFF);
+			words.add(new String(chunk,0, inputStream.read(chunk), StandardCharsets.UTF_8));
+			assertThat(inputStream.read()).isEqualTo(' ' & 0xFF);
+			words.add(new String(chunk,0, inputStream.read(chunk), StandardCharsets.UTF_8));
+			assertThat(inputStream.read()).isEqualTo(-1);
+		}
+		catch (IOException e) {
+			throw new RuntimeException(e);
+		}
+		assertThat(words).containsExactlyElementsOf(List.of("foo", "bar", "baz"));
+	}
+
+	@ParameterizedDataBufferAllocatingTest
+	void inputStreamSubscriberClose(DataBufferFactory bufferFactory) throws InterruptedException {
+		for (int i = 1; i < 100; i++) {
+			CountDownLatch latch = new CountDownLatch(1);
+			super.bufferFactory = bufferFactory;
+
+			var input = List.of("foo", "bar", "baz");
+
+			Publisher<DataBuffer> publisher = DataBufferUtils.outputStreamPublisher(outputStream -> {
+				try {
+					assertThatIOException()
+							.isThrownBy(() -> {
+								for (String word : input) {
+									outputStream.write(word.getBytes(StandardCharsets.UTF_8));
+									outputStream.flush();
+								}
+							})
+							.withMessage("Subscription has been terminated");
+				} finally {
+					latch.countDown();
+				}
+			}, super.bufferFactory, Executors.newSingleThreadExecutor(), 1);
+
+
+			byte[] chunk = new byte[3];
+			ArrayList<String> words = new ArrayList<>();
+
+			try (InputStream inputStream = DataBufferUtils.subscribeAsInputStream(publisher, ThreadLocalRandom.current().nextInt(1, 4))) {
+				inputStream.read(chunk);
+				String word = new String(chunk, StandardCharsets.UTF_8);
+				words.add(word);
+			} catch (IOException e) {
+				throw new RuntimeException(e);
+			}
+			assertThat(words).containsExactlyElementsOf(List.of("foo"));
+			latch.await();
+		}
+	}
+
 	@ParameterizedDataBufferAllocatingTest
 	void readAndWriteByteChannel(DataBufferFactory bufferFactory) throws Exception {
 		super.bufferFactory = bufferFactory;