From d626b5e833ecd76f55e9118060a4e7d523b11971 Mon Sep 17 00:00:00 2001 From: Arjen Poutsma Date: Wed, 20 Apr 2016 13:30:40 +0200 Subject: [PATCH] Various DataBuffer improvements This commit introduces two DataBuffer improvements: - The capability to read a Flux from an input stream or channel. - The capability to limit a Publisher to publish up until a given maximum byte count. --- .../core/io/buffer/DefaultDataBuffer.java | 6 +- .../io/buffer/support/DataBufferUtils.java | 226 ++++++++++++++++-- .../buffer/support/DataBufferUtilsTests.java | 123 ++++++++++ .../buffer/support/DataBufferUtilsTests.txt | 4 + 4 files changed, 334 insertions(+), 25 deletions(-) create mode 100644 spring-web-reactive/src/test/java/org/springframework/core/io/buffer/support/DataBufferUtilsTests.java create mode 100644 spring-web-reactive/src/test/resources/org/springframework/core/io/buffer/support/DataBufferUtilsTests.txt diff --git a/spring-web-reactive/src/main/java/org/springframework/core/io/buffer/DefaultDataBuffer.java b/spring-web-reactive/src/main/java/org/springframework/core/io/buffer/DefaultDataBuffer.java index 0dca055bec..823f644eec 100644 --- a/spring-web-reactive/src/main/java/org/springframework/core/io/buffer/DefaultDataBuffer.java +++ b/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 writeInternal(Function 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; } diff --git a/spring-web-reactive/src/main/java/org/springframework/core/io/buffer/support/DataBufferUtils.java b/spring-web-reactive/src/main/java/org/springframework/core/io/buffer/support/DataBufferUtils.java index 745b42d42e..16fd7e5b8e 100644 --- a/spring-web-reactive/src/main/java/org/springframework/core/io/buffer/support/DataBufferUtils.java +++ b/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 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 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 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 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 Consumer closeConsumer() { + return (Consumer) 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 takeUntilByteCount(Publisher 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() { + + 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> { + + 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 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> { + + 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 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); + } + + } + } + } diff --git a/spring-web-reactive/src/test/java/org/springframework/core/io/buffer/support/DataBufferUtilsTests.java b/spring-web-reactive/src/test/java/org/springframework/core/io/buffer/support/DataBufferUtilsTests.java new file mode 100644 index 0000000000..052c6a2a00 --- /dev/null +++ b/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 flux = DataBufferUtils.read(channel, allocator, 4); + + TestSubscriber 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 flux = DataBufferUtils.read(channel, allocator, 3); + + TestSubscriber 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 flux = DataBufferUtils.read(is, allocator, 4); + + TestSubscriber 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 flux = DataBufferUtils.read(is, allocator, 3); + + TestSubscriber 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 flux = + Flux.just(stringBuffer("foo"), stringBuffer("bar"), stringBuffer("baz")); + + Flux result = DataBufferUtils.takeUntilByteCount(flux, 5L); + + TestSubscriber testSubscriber = new TestSubscriber<>(); + testSubscriber.bindTo(result). + assertNoError(). + assertComplete(). + assertValues(stringBuffer("foo"), stringBuffer("ba")); + } + + +} \ No newline at end of file diff --git a/spring-web-reactive/src/test/resources/org/springframework/core/io/buffer/support/DataBufferUtilsTests.txt b/spring-web-reactive/src/test/resources/org/springframework/core/io/buffer/support/DataBufferUtilsTests.txt new file mode 100644 index 0000000000..ab9b661144 --- /dev/null +++ b/spring-web-reactive/src/test/resources/org/springframework/core/io/buffer/support/DataBufferUtilsTests.txt @@ -0,0 +1,4 @@ +foo +bar +baz +qux