Add optimized DataBufferInputStream overrides

Add optimized DataBufferInputStream overrides for readNBytes, skip, and transferTo; all of them
allocate byte buffers which we can either avoid (in the case of skip) or size more precisely since
the number of remaining bytes is known.

Signed-off-by: Patrick Strawderman <[email protected]>

Author: kilink

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

@@ -18,6 +18,11 @@
 
 import java.io.IOException;
 import java.io.InputStream;
+import java.io.OutputStream;
+import java.nio.ByteBuffer;
+import java.nio.channels.Channels;
+import java.nio.channels.WritableByteChannel;
+import java.util.Objects;
 
 import org.springframework.util.Assert;
 
@@ -103,10 +108,44 @@ public void close() {
 		this.closed = true;
 	}
 
+	@Override
+	public byte[] readNBytes(int len) throws IOException {
+		if (len < 0) {
+			throw new IllegalArgumentException("len < 0");
+		}
+		checkClosed();
+		int size = Math.min(available(), len);
+		byte[] out = new byte[size];
+		this.dataBuffer.read(out);
+		return out;
+	}
+
+	@Override
+	public long skip(long n) throws IOException {
+		checkClosed();
+		if (n <= 0) {
+			return 0L;
+		}
+		int skipped = Math.min(available(), n > Integer.MAX_VALUE ? Integer.MAX_VALUE : (int) n);
+		this.dataBuffer.readPosition(Math.min(this.end, this.dataBuffer.readPosition() + skipped));
+		return skipped;
+	}
+
+	@Override
+	public long transferTo(OutputStream out) throws IOException {
+		Objects.requireNonNull(out, "out");
+		checkClosed();
+		if (available() == 0) {
+			return 0L;
+		}
+		byte[] buf = readAllBytes();
+		out.write(buf);
+		return buf.length;
+	}
+
 	private void checkClosed() throws IOException {
 		if (this.closed) {
 			throw new IOException("DataBufferInputStream is closed");
 		}
 	}
-
 }

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

@@ -16,6 +16,7 @@
 
 package org.springframework.core.io.buffer;
 
+import java.io.ByteArrayOutputStream;
 import java.io.InputStream;
 import java.io.OutputStream;
 import java.nio.ByteBuffer;
@@ -341,6 +342,48 @@ void inputStream(DataBufferFactory bufferFactory) throws Exception {
 		assertThat(len).isEqualTo(3);
 		assertThat(bytes).containsExactly('c', 'd', 'e');
 
+		buffer.readPosition(0);
+		inputStream = buffer.asInputStream();
+		assertThat(inputStream.readAllBytes()).asString().isEqualTo("abcde");
+		assertThat(inputStream.available()).isEqualTo(0);
+		assertThat(inputStream.readAllBytes()).isEmpty();
+
+		buffer.readPosition(0);
+		inputStream = buffer.asInputStream();
+		inputStream.mark(5);
+		assertThat(inputStream.readNBytes(0)).isEmpty();
+		assertThat(inputStream.readNBytes(1000)).asString().isEqualTo("abcde");
+		inputStream.reset();
+		assertThat(inputStream.readNBytes(3)).asString().isEqualTo("abc");
+		assertThat(inputStream.readNBytes(2)).asString().isEqualTo("de");
+		assertThat(inputStream.readNBytes(10)).isEmpty();
+
+		buffer.readPosition(0);
+		inputStream = buffer.asInputStream();
+		inputStream.mark(5);
+		assertThat(inputStream.skip(1)).isEqualTo(1);
+		assertThat(inputStream.readAllBytes()).asString().isEqualTo("bcde");
+		assertThat(inputStream.skip(10)).isEqualTo(0);
+		assertThat(inputStream.available()).isEqualTo(0);
+		inputStream.reset();
+		assertThat(inputStream.skip(100)).isEqualTo(5);
+		assertThat(inputStream.available()).isEqualTo(0);
+
+		buffer.readPosition(0);
+		inputStream = buffer.asInputStream();
+		inputStream.mark(5);
+		ByteArrayOutputStream out = new ByteArrayOutputStream();
+		assertThat(inputStream.transferTo(out)).isEqualTo(5);
+		assertThat(out.toByteArray()).asString().isEqualTo("abcde");
+		assertThat(inputStream.available()).isEqualTo(0);
+		out.reset();
+		inputStream.reset();
+		assertThat(inputStream.read()).isEqualTo('a');
+		assertThat(inputStream.transferTo(out)).isEqualTo(4);
+		assertThat(out.toByteArray()).asString().isEqualTo("bcde");
+		assertThat(inputStream.available()).isEqualTo(0);
+		assertThat(inputStream.transferTo(OutputStream.nullOutputStream())).isEqualTo(0);
+
 		release(buffer);
 	}