Various DataBuffer improvements

This commit introduces two DataBuffer improvements:

 - The capability to read a Flux<DataBuffer> from an input stream or
   channel.

 - The capability to limit a Publisher<DataBuffer> to publish up until a
   given maximum byte count.

spring-web-reactive/src/main/java/org/springframework/core/io/buffer/DefaultDataBuffer.java

@@ -179,7 +179,7 @@ public class DefaultDataBuffer implements DataBuffer {
 	}
 
 	/**
-	 * Internal write method that keeps track of the {@link #writePosition} befor eand
+	 * Internal write method that keeps track of the {@link #writePosition} before and
 	 * after applying the given function on {@link #byteBuffer}.
 	 */
 	private <T> T writeInternal(Function<ByteBuffer, T> function) {
@@ -239,7 +239,9 @@ public class DefaultDataBuffer implements DataBuffer {
 		}
 		else if (obj instanceof DefaultDataBuffer) {
 			DefaultDataBuffer other = (DefaultDataBuffer) obj;
-			return this.byteBuffer.equals(other.byteBuffer);
+			return this.readPosition == other.readPosition &&
+					this.writePosition == other.writePosition &&
+					this.byteBuffer.equals(other.byteBuffer);
 		}
 		return false;
 	}

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

@@ -16,47 +16,44 @@
 
 package org.springframework.core.io.buffer.support;
 
+import java.io.Closeable;
+import java.io.IOException;
 import java.io.InputStream;
 import java.io.SequenceInputStream;
+import java.nio.ByteBuffer;
+import java.nio.channels.ReadableByteChannel;
 import java.util.Enumeration;
 import java.util.Iterator;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
 
 import org.reactivestreams.Publisher;
+import org.reactivestreams.Subscriber;
+import org.reactivestreams.Subscription;
 import reactor.core.publisher.Flux;
+import reactor.core.subscriber.SubscriberWithContext;
 
 import org.springframework.core.io.buffer.DataBuffer;
+import org.springframework.core.io.buffer.DataBufferAllocator;
 import org.springframework.util.Assert;
 
-/**
+/**i
  * Utility class for working with {@link DataBuffer}s.
  *
  * @author Arjen Poutsma
  */
 public abstract class DataBufferUtils {
 
-	/**
-	 * Returns the given {@link DataBuffer} as a {@link Flux} of bytes.
-	 * @param buffer the buffer to return the bytes of
-	 * @return the bytes as a flux
-	 */
-	public static Flux<Byte> toPublisher(DataBuffer buffer) {
-		Assert.notNull(buffer, "'buffer' must not be null");
-
-		byte[] bytes = new byte[buffer.readableByteCount()];
-		buffer.read(bytes);
-
-		Byte[] bytesObjects = box(bytes);
-
-		return Flux.fromArray(bytesObjects);
-	}
-
-	private static Byte[] box(byte[] bytes) {
-		Byte[] bytesObjects = new Byte[bytes.length];
-		for (int i = 0; i < bytes.length; i++) {
-			bytesObjects[i] = bytes[i];
+	private static final Consumer<? extends Closeable> CLOSE_CONSUMER = closeable -> {
+		try {
+			if (closeable != null) {
+				closeable.close();
+			}
 		}
-		return bytesObjects;
-	}
+		catch (IOException ignored) {
+		}
+	};
+
 
 	/**
 	 * Returns the given data buffer publisher as a blocking input stream, streaming over
@@ -75,6 +72,103 @@ public abstract class DataBufferUtils {
 		return new SequenceInputStream(enumeration);
 	}
 
+	/**
+	 * Reads the given {@code ReadableByteChannel} into a {@code Flux} of
+	 * {@code DataBuffer}s. Closes the channel when the flux is terminated.
+	 * @param channel the channel to read from
+	 * @param allocator the allocator to create data buffers with
+	 * @param bufferSize the maximum size of the data buffers
+	 * @return a flux of data buffers read from the given channel
+	 */
+	public static Flux<DataBuffer> read(ReadableByteChannel channel,
+			DataBufferAllocator allocator, int bufferSize) {
+		Assert.notNull(channel, "'channel' must not be null");
+		Assert.notNull(allocator, "'allocator' must not be null");
+
+		return Flux.create(new ReadableByteChannelConsumer(allocator, bufferSize),
+				subscriber -> channel, closeConsumer());
+	}
+
+	/**
+	 * Reads the given {@code InputStream} into a {@code Flux} of
+	 * {@code DataBuffer}s. Closes the stream when the flux inputStream terminated.
+	 * @param inputStream the input stream to read from
+	 * @param allocator the allocator to create data buffers with
+	 * @param bufferSize the maximum size of the data buffers
+	 * @return a flux of data buffers read from the given channel
+	 */
+	public static Flux<DataBuffer> read(InputStream inputStream,
+			DataBufferAllocator allocator, int bufferSize) {
+		Assert.notNull(inputStream, "'inputStream' must not be null");
+		Assert.notNull(allocator, "'allocator' must not be null");
+
+		return Flux.create(new InputStreamConsumer(allocator, bufferSize),
+				subscriber -> inputStream, closeConsumer());
+	}
+
+	@SuppressWarnings("unchecked")
+	private static <T extends Closeable> Consumer<T> closeConsumer() {
+		return (Consumer<T>) CLOSE_CONSUMER;
+	}
+
+	/**
+	 * Relays buffers from the given {@link Publisher} until the total
+	 * {@linkplain DataBuffer#readableByteCount() byte count} reaches the given maximum
+	 * byte count, or until the publisher is complete.
+	 * @param publisher the publisher to filter
+	 * @param maxByteCount the maximum byte count
+	 * @return a flux whose maximum byte count is {@code maxByteCount}
+	 */
+	public static Flux<DataBuffer> takeUntilByteCount(Publisher<DataBuffer> publisher,
+			long maxByteCount) {
+		Assert.notNull(publisher, "'publisher' must not be null");
+		Assert.isTrue(maxByteCount >= 0, "'maxByteCount' must be a positive number");
+
+		return Flux.from(publisher).lift(subscriber -> new Subscriber<DataBuffer>() {
+
+			private Subscription subscription;
+
+			private final AtomicLong byteCount = new AtomicLong();
+
+			@Override
+			public void onSubscribe(Subscription s) {
+				this.subscription = s;
+				subscriber.onSubscribe(s);
+			}
+
+			@Override
+			public void onNext(DataBuffer dataBuffer) {
+				long currentCount =
+						this.byteCount.addAndGet(dataBuffer.readableByteCount());
+				if (currentCount > maxByteCount) {
+					int size = (int) (currentCount - maxByteCount + 1);
+					ByteBuffer byteBuffer =
+							(ByteBuffer) dataBuffer.asByteBuffer().limit(size);
+					DataBuffer partialBuffer =
+							dataBuffer.allocator().allocateBuffer(size);
+					partialBuffer.write(byteBuffer);
+
+					subscriber.onNext(partialBuffer);
+					subscriber.onComplete();
+					this.subscription.cancel();
+				}
+				else {
+					subscriber.onNext(dataBuffer);
+				}
+			}
+
+			@Override
+			public void onError(Throwable t) {
+				subscriber.onError(t);
+			}
+
+			@Override
+			public void onComplete() {
+				subscriber.onComplete();
+			}
+		});
+	}
+
 	/**
 	 * Enumeration wrapping an Iterator.
 	 */
@@ -98,4 +192,90 @@ public abstract class DataBufferUtils {
 		}
 	}
 
+	private static class ReadableByteChannelConsumer
+			implements Consumer<SubscriberWithContext<DataBuffer, ReadableByteChannel>> {
+
+		private final DataBufferAllocator allocator;
+
+		private final int chunkSize;
+
+		public ReadableByteChannelConsumer(DataBufferAllocator allocator, int chunkSize) {
+			this.allocator = allocator;
+			this.chunkSize = chunkSize;
+		}
+
+		@Override
+		public void accept(SubscriberWithContext<DataBuffer, ReadableByteChannel> sub) {
+			try {
+				ByteBuffer byteBuffer = ByteBuffer.allocate(chunkSize);
+				int read;
+				ReadableByteChannel channel = sub.context();
+				if ((read = channel.read(byteBuffer)) > 0) {
+					byteBuffer.flip();
+					boolean release = true;
+					DataBuffer dataBuffer = this.allocator.allocateBuffer(read);
+					try {
+						dataBuffer.write(byteBuffer);
+						release = false;
+						sub.onNext(dataBuffer);
+					}
+					finally {
+						if (release) {
+							// TODO: release buffer when we have PooledDataBuffer
+						}
+					}
+				}
+				else {
+					sub.onComplete();
+				}
+			}
+			catch (IOException ex) {
+				sub.onError(ex);
+			}
+		}
+	}
+
+	private static class InputStreamConsumer
+			implements Consumer<SubscriberWithContext<DataBuffer, InputStream>> {
+
+		private final DataBufferAllocator allocator;
+
+		private final int chunkSize;
+
+		public InputStreamConsumer(DataBufferAllocator allocator, int chunkSize) {
+			this.allocator = allocator;
+			this.chunkSize = chunkSize;
+		}
+
+		@Override
+		public void accept(SubscriberWithContext<DataBuffer, InputStream> sub) {
+			try {
+				byte[] bytes = new byte[chunkSize];
+				int read;
+				InputStream is = sub.context();
+				if ((read = is.read(bytes)) > 0) {
+					boolean release = true;
+					DataBuffer dataBuffer = this.allocator.allocateBuffer(read);
+					try {
+						dataBuffer.write(bytes, 0, read);
+						release = false;
+						sub.onNext(dataBuffer);
+					}
+					finally {
+						if (release) {
+							// TODO: release buffer when we have PooledDataBuffer
+						}
+					}
+				}
+				else {
+					sub.onComplete();
+				}
+			}
+			catch (IOException ex) {
+				sub.onError(ex);
+			}
+
+		}
+	}
+
 }

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

@@ -0,0 +1,123 @@
+/*
+ * Copyright 2002-2016 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.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.core.io.buffer.support;
+
+import java.io.InputStream;
+import java.net.URI;
+import java.nio.channels.FileChannel;
+import java.nio.file.Paths;
+import java.nio.file.StandardOpenOption;
+
+import org.junit.Test;
+import reactor.core.publisher.Flux;
+import reactor.core.test.TestSubscriber;
+
+import org.springframework.core.codec.support.AbstractAllocatingTestCase;
+import org.springframework.core.io.buffer.DataBuffer;
+
+import static org.junit.Assert.assertFalse;
+
+/**
+ * @author Arjen Poutsma
+ */
+public class DataBufferUtilsTests extends AbstractAllocatingTestCase {
+
+	@Test
+	public void readChannel() throws Exception {
+		URI uri = DataBufferUtilsTests.class.getResource("DataBufferUtilsTests.txt")
+				.toURI();
+		FileChannel channel = FileChannel.open(Paths.get(uri), StandardOpenOption.READ);
+
+		Flux<DataBuffer> flux = DataBufferUtils.read(channel, allocator, 4);
+
+		TestSubscriber<DataBuffer> testSubscriber = new TestSubscriber<>();
+		testSubscriber.bindTo(flux).
+				assertNoError().
+				assertComplete().
+				assertValues(stringBuffer("foo\n"), stringBuffer("bar\n"),
+						stringBuffer("baz\n"), stringBuffer("qux\n"));
+
+		assertFalse(channel.isOpen());
+	}
+
+	@Test
+	public void readUnalignedChannel() throws Exception {
+		URI uri = DataBufferUtilsTests.class.getResource("DataBufferUtilsTests.txt")
+				.toURI();
+		FileChannel channel = FileChannel.open(Paths.get(uri), StandardOpenOption.READ);
+
+		Flux<DataBuffer> flux = DataBufferUtils.read(channel, allocator, 3);
+
+		TestSubscriber<DataBuffer> testSubscriber = new TestSubscriber<>();
+		testSubscriber.bindTo(flux).
+				assertNoError().
+				assertComplete().
+				assertValues(stringBuffer("foo"), stringBuffer("\nba"),
+						stringBuffer("r\nb"), stringBuffer("az\n"), stringBuffer("qux"),
+						stringBuffer("\n"));
+
+		assertFalse(channel.isOpen());
+	}
+
+	@Test
+	public void readInputStream() {
+		InputStream is = DataBufferUtilsTests.class
+				.getResourceAsStream("DataBufferUtilsTests.txt");
+
+		Flux<DataBuffer> flux = DataBufferUtils.read(is, allocator, 4);
+
+		TestSubscriber<DataBuffer> testSubscriber = new TestSubscriber<>();
+		testSubscriber.bindTo(flux).
+				assertNoError().
+				assertComplete().
+				assertValues(stringBuffer("foo\n"), stringBuffer("bar\n"),
+						stringBuffer("baz\n"), stringBuffer("qux\n"));
+	}
+
+	@Test
+	public void readUnalignedInputStream() throws Exception {
+		InputStream is = DataBufferUtilsTests.class
+				.getResourceAsStream("DataBufferUtilsTests.txt");
+
+		Flux<DataBuffer> flux = DataBufferUtils.read(is, allocator, 3);
+
+		TestSubscriber<DataBuffer> testSubscriber = new TestSubscriber<>();
+		testSubscriber.bindTo(flux).
+				assertNoError().
+				assertComplete().
+				assertValues(stringBuffer("foo"), stringBuffer("\nba"),
+						stringBuffer("r\nb"), stringBuffer("az\n"), stringBuffer("qux"),
+						stringBuffer("\n"));
+	}
+
+
+	@Test
+	public void takeUntilByteCount() {
+		Flux<DataBuffer> flux =
+				Flux.just(stringBuffer("foo"), stringBuffer("bar"), stringBuffer("baz"));
+
+		Flux<DataBuffer> result = DataBufferUtils.takeUntilByteCount(flux, 5L);
+
+		TestSubscriber<DataBuffer> testSubscriber = new TestSubscriber<>();
+		testSubscriber.bindTo(result).
+				assertNoError().
+				assertComplete().
+				assertValues(stringBuffer("foo"), stringBuffer("ba"));
+	}
+
+
+}

spring-web-reactive/src/test/resources/org/springframework/core/io/buffer/support/DataBufferUtilsTests.txt

@@ -0,0 +1,4 @@
+foo
+bar
+baz
+qux