diff --git a/src/main/java/org/elasticsearch/common/io/stream/BytesStreamOutput.java b/src/main/java/org/elasticsearch/common/io/stream/BytesStreamOutput.java index 378a5fecb7a..52756377b5a 100644 --- a/src/main/java/org/elasticsearch/common/io/stream/BytesStreamOutput.java +++ b/src/main/java/org/elasticsearch/common/io/stream/BytesStreamOutput.java @@ -31,7 +31,9 @@ import java.io.IOException; */ public class BytesStreamOutput extends StreamOutput implements BytesStream { - public static final int DEFAULT_SIZE = 32 * 1024; + public static final int DEFAULT_SIZE = 2 * 1024; + + public static final int OVERSIZE_LIMIT = 256 * 1024; /** * The buffer where data is stored. @@ -73,7 +75,7 @@ public class BytesStreamOutput extends StreamOutput implements BytesStream { public void writeByte(byte b) throws IOException { int newcount = count + 1; if (newcount > buf.length) { - buf = ArrayUtil.grow(buf, newcount); + buf = grow(newcount); } buf[count] = b; count = newcount; @@ -82,7 +84,7 @@ public class BytesStreamOutput extends StreamOutput implements BytesStream { public void skip(int length) { int newcount = count + length; if (newcount > buf.length) { - buf = ArrayUtil.grow(buf, newcount); + buf = grow(newcount); } count = newcount; } @@ -94,12 +96,20 @@ public class BytesStreamOutput extends StreamOutput implements BytesStream { } int newcount = count + length; if (newcount > buf.length) { - buf = ArrayUtil.grow(buf, newcount); + buf = grow(newcount); } System.arraycopy(b, offset, buf, count, length); count = newcount; } + private byte[] grow(int newCount) { + // try and grow faster while we are small... + if (newCount < OVERSIZE_LIMIT) { + newCount = Math.max(buf.length << 1, newCount); + } + return ArrayUtil.grow(buf, newCount); + } + public void seek(int seekTo) { count = seekTo; } @@ -108,6 +118,10 @@ public class BytesStreamOutput extends StreamOutput implements BytesStream { count = 0; } + public int bufferSize() { + return buf.length; + } + @Override public void flush() throws IOException { // nothing to do there diff --git a/src/test/java/org/elasticsearch/test/unit/common/io/streams/BytesStreamsTests.java b/src/test/java/org/elasticsearch/test/unit/common/io/streams/BytesStreamsTests.java index c1f4dc450e3..d4ee4fff535 100644 --- a/src/test/java/org/elasticsearch/test/unit/common/io/streams/BytesStreamsTests.java +++ b/src/test/java/org/elasticsearch/test/unit/common/io/streams/BytesStreamsTests.java @@ -60,4 +60,17 @@ public class BytesStreamsTests { assertThat(in.readString(), equalTo("hello")); assertThat(in.readString(), equalTo("goodbye")); } + + @Test + public void testGrowLogic() throws Exception { + BytesStreamOutput out = new BytesStreamOutput(); + out.writeBytes(new byte[BytesStreamOutput.DEFAULT_SIZE - 5]); + assertThat(out.bufferSize(), equalTo(2048)); // remains the default + out.writeBytes(new byte[1 * 1024]); + assertThat(out.bufferSize(), equalTo(4608)); + out.writeBytes(new byte[32 * 1024]); + assertThat(out.bufferSize(), equalTo(40320)); + out.writeBytes(new byte[32 * 1024]); + assertThat(out.bufferSize(), equalTo(90720)); + } }