Fix off-by-one bug in RecyclerBytesStreamOutput (#95036) (#95043)

Fixing off by one bug when seeking to first byte in the next page.
elastic · Apr 5, 2023 · 92239bd · 92239bd
1 parent f7baae9
commit 92239bd
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 4 deletions.
diff --git a/docs/changelog/95036.yaml b/docs/changelog/95036.yaml
@@ -0,0 +1,5 @@
+pr: 95036
+summary: Fix off-by-one bug in `RecyclerBytesStreamOutput`
+area: Network
+type: bug
+issues: []
diff --git a/server/src/main/java/org/elasticsearch/common/io/stream/RecyclerBytesStreamOutput.java b/server/src/main/java/org/elasticsearch/common/io/stream/RecyclerBytesStreamOutput.java
@@ -230,7 +230,7 @@ private void ensureCapacityFromPosition(long newPosition) {
         if (newPosition > Integer.MAX_VALUE - (Integer.MAX_VALUE % pageSize)) {
             throw new IllegalArgumentException(getClass().getSimpleName() + " cannot hold more than 2GB of data");
         }
-        while (newPosition > currentCapacity) {
+        while (newPosition > currentCapacity - 1) {
             Recycler.V<BytesRef> newPage = recycler.obtain();
             assert pageSize == newPage.v().length;
             pages.add(newPage);

diff --git a/server/src/test/java/org/elasticsearch/common/io/stream/RecyclerBytesStreamOutputTests.java b/server/src/test/java/org/elasticsearch/common/io/stream/RecyclerBytesStreamOutputTests.java
@@ -20,6 +20,7 @@
 import org.elasticsearch.common.util.Maps;
 import org.elasticsearch.common.util.PageCacheRecycler;
 import org.elasticsearch.core.TimeValue;
+import org.elasticsearch.core.Tuple;
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.transport.BytesRefRecycler;
 
@@ -240,12 +241,43 @@ public void testSeek() throws Exception {
         int position = 0;
         assertEquals(position, out.position());
 
-        out.seek(position += 10);
+        final List<Tuple<Integer, Byte>> writes = new ArrayList<>();
+        int randomOffset = randomIntBetween(0, PageCacheRecycler.BYTE_PAGE_SIZE);
+        out.seek(position += randomOffset);
+        if (randomBoolean()) {
+            byte random = randomByte();
+            out.writeByte(random);
+            writes.add(Tuple.tuple(position, random));
+            position++;
+        }
         out.seek(position += PageCacheRecycler.BYTE_PAGE_SIZE);
-        out.seek(position += PageCacheRecycler.BYTE_PAGE_SIZE + 10);
+        if (randomBoolean()) {
+            byte random = randomByte();
+            out.writeByte(random);
+            writes.add(Tuple.tuple(position, random));
+            position++;
+        }
+        out.seek(position += PageCacheRecycler.BYTE_PAGE_SIZE + randomOffset);
+        if (randomBoolean()) {
+            byte random = randomByte();
+            out.writeByte(random);
+            writes.add(Tuple.tuple(position, random));
+            position++;
+        }
         out.seek(position += PageCacheRecycler.BYTE_PAGE_SIZE * 2);
+        if (randomBoolean()) {
+            byte random = randomByte();
+            out.writeByte(random);
+            writes.add(Tuple.tuple(position, random));
+            position++;
+        }
         assertEquals(position, out.position());
-        assertEquals(position, BytesReference.toBytes(out.bytes()).length);
+
+        final BytesReference bytesReference = out.bytes();
+        assertEquals(position, BytesReference.toBytes(bytesReference).length);
+        for (Tuple<Integer, Byte> write : writes) {
+            assertEquals((byte) write.v2(), bytesReference.get(write.v1()));
+        }
 
         IllegalArgumentException iae = expectThrows(IllegalArgumentException.class, () -> out.seek(Integer.MAX_VALUE + 1L));
         assertEquals("RecyclerBytesStreamOutput cannot hold more than 2GB of data", iae.getMessage());
@@ -1002,4 +1034,12 @@ public V<BytesRef> obtain() {
             assertThat(pagesAllocated.get(), equalTo(0L));
         }
     }
+
+    public void testSeekToPageBoundary() {
+        RecyclerBytesStreamOutput out = new RecyclerBytesStreamOutput(recycler);
+        out.seek(PageCacheRecycler.BYTE_PAGE_SIZE);
+        byte b = randomByte();
+        out.writeByte(b);
+        assertEquals(b, out.bytes().get(PageCacheRecycler.BYTE_PAGE_SIZE));
+    }
 }