Repository: amzn/ion-java
Pull request #416 (2022-03-01 19:38:08)

Fixes a bug in WriteBuffer#shiftBytesLeft
This patch fixes a bug in the `WriteBuffer` that caused the
`shiftBytesLeft` method to corrupt data if it was called such that the
final `Block` in the list was emptied AND the second-to-last block
reclaimed some capacity. This would only occur if the writer
was configured to use length preallocation greater than 0 and if they
were writing more than 32KB of Ion (the default block size) between
flushes.

This change accounts for that case and, if the last `Block`
is emptied, returns it to the pool for reuse.

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.


tgregg: I think the following test should pass, but it fails with `ArrayIndexOutOfBoundsException` on the line `buf.shiftBytesLeft(1, 6)`, for any shift amount >= 6 (which is just across the block boundary). It works for shift amounts of 5 or less (remaining in the same block).

```java
    @Test
    public void shiftBytesLeftAcrossBufferBlocksToStartThenWritePastBlockBoundary() throws IOException {
        assertEquals(11, ALLOCATOR.getBlockSize());
        // The "B" is the first and only byte in the second block.
        // Shifting 2 bytes left by 10 empties the last block and positions the shifted bytes at the start of the first
        // block.
        buf.writeBytes("0123456789AB".getBytes());
        buf.shiftBytesLeft(2, 10);
        assertBuffer("AB".getBytes());
        // Write enough bytes to expand back into the next block.
        buf.writeBytes("CDE0123456789AB".getBytes());
        assertBuffer("ABCDE0123456789AB".getBytes());
        // Shift one byte left past the block boundary. This should once again empty the last block.
        buf.shiftBytesLeft(1, 6);
        assertBuffer("ABCDE01234B".getBytes());
    }
```

@zslayton can you double check my test logic?
zslayton: > 



> I think the following test should pass, but it fails with `ArrayIndexOutOfBoundsException` on the line `buf.shiftBytesLeft(1, 6)`, for any shift amount >= 6 (which is just across the block boundary). It works for shift amounts of 5 or less (remaining in the same block).
> 
> ```java
>     @Test
>     public void shiftBytesLeftAcrossBufferBlocksToStartThenWritePastBlockBoundary() throws IOException {
>         assertEquals(11, ALLOCATOR.getBlockSize());
>         // The "B" is the first and only byte in the second block.
>         // Shifting 2 bytes left by 10 empties the last block and positions the shifted bytes at the start of the first
>         // block.
>         buf.writeBytes("0123456789AB".getBytes());
>         buf.shiftBytesLeft(2, 10);
>         assertBuffer("AB".getBytes());
>         // Write enough bytes to expand back into the next block.
>         buf.writeBytes("CDE0123456789AB".getBytes());
>         assertBuffer("ABCDE0123456789AB".getBytes());
>         // Shift one byte left past the block boundary. This should once again empty the last block.
>         buf.shiftBytesLeft(1, 6);
>         assertBuffer("ABCDE01234B".getBytes());
>     }
> ```
> 
> @zslayton can you double check my test logic?

`shiftBytesLeft` is intended to shift the last `N` bytes to the left by up to `N` positions. (In practice, it's only called when `length >= 14` and `shiftBy <= 2`.) It's not clear that shifting `N` bytes by more than `N` positions should be a supported use case.

If you're alright with it, I think I'll address this by checking for this in `shiftBytesLeft` and throwing an `IllegalArgumentException` if it arises. I'll also note it in the docs and add a unit test for it.
tgregg: > If you're alright with it, I think I'll address this by checking for this in shiftBytesLeft and throwing an IllegalArgumentException if it arises. I'll also note it in the docs and add a unit test for it.

Yep, that works for me.
codecov[bot]: # [Codecov](https://codecov.io/gh/amzn/ion-java/pull/416?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) Report
> Merging [#416](https://codecov.io/gh/amzn/ion-java/pull/416?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) (fd29a11) into [master](https://codecov.io/gh/amzn/ion-java/commit/3efa93a44479890e14e48501160a69b53a872af4?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) (3efa93a) will **increase** coverage by `0.03%`.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/amzn/ion-java/pull/416/graphs/tree.svg?width=650&height=150&src=pr&token=M93Y3e3kKn&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None)](https://codecov.io/gh/amzn/ion-java/pull/416?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None)

```diff
@@             Coverage Diff              @@
##             master     #416      +/-   ##
============================================
+ Coverage     66.45%   66.48%   +0.03%     
- Complexity     5402     5405       +3     
============================================
  Files           156      156              
  Lines         22745    22760      +15     
  Branches       4099     4102       +3     
============================================
+ Hits          15116    15133      +17     
+ Misses         6266     6265       -1     
+ Partials       1363     1362       -1     
```


| [Impacted Files](https://codecov.io/gh/amzn/ion-java/pull/416?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) | Coverage Δ | |
|---|---|---|
| [src/com/amazon/ion/impl/bin/WriteBuffer.java](https://codecov.io/gh/amzn/ion-java/pull/416/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None#diff-c3JjL2NvbS9hbWF6b24vaW9uL2ltcGwvYmluL1dyaXRlQnVmZmVyLmphdmE=) | `94.40% <100.00%> (+0.15%)` | :arrow_up: |
| [src/com/amazon/ion/impl/BlockedBuffer.java](https://codecov.io/gh/amzn/ion-java/pull/416/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None#diff-c3JjL2NvbS9hbWF6b24vaW9uL2ltcGwvQmxvY2tlZEJ1ZmZlci5qYXZh) | `51.45% <0.00%> (+0.24%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/amzn/ion-java/pull/416?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/amzn/ion-java/pull/416?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None). Last update [3efa93a...fd29a11](https://codecov.io/gh/amzn/ion-java/pull/416?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None).

zslayton: I rewrote `shiftBytesLeftAcrossBlocks` from scratch and I think the updated version is _much_ easier to follow. The new version has far fewer branches and eliminates a lot of the restrictions that were in the original version. I've also added several new unit tests, including the one Tyler provided.

This last commit obviates a lot of the changes in the first commit. Instead of viewing "Changes since my last review", I suggest looking at all the commits together; it's more straightforward.

Patch:
diff --git a/src/com/amazon/ion/impl/bin/WriteBuffer.java b/src/com/amazon/ion/impl/bin/WriteBuffer.java
index 0df5ba68..c4550dc3 100644
--- a/src/com/amazon/ion/impl/bin/WriteBuffer.java
+++ b/src/com/amazon/ion/impl/bin/WriteBuffer.java
@@ -180,7 +180,7 @@ import java.util.List;
      * @param shiftBy   The number of bytes to the left that we'll be shifting.
      */
     public void shiftBytesLeft(int length, int shiftBy) {
-        if (shiftBy == 0) {
+        if (length == 0 || shiftBy == 0) {
             // Nothing to do.
             return;
         }
@@ -227,63 +227,74 @@ import java.util.List;
      * @param shiftBy   The number of bytes to the left that we'll be shifting.
      */
     private void shiftBytesLeftAcrossBlocks(int length, int shiftBy) {
-        // Our starting position is the first byte that we plan to shift backwards. The `bufferOffset` is the cursor's
-        // current position in the stream; we can use this to derive which block it's in as well as its position within
-        // that block.
-        long bufferOffset = position() - length;
-
-        Block block = null;
-        while (length > 0)
-        {
-            // Using the `bufferOffset`, determine which block we're in...
-            int blockIndex = index(bufferOffset);
-            block = blocks.get(blockIndex);
-            // ...and our offset within that block.
-            int blockOffset = offset(bufferOffset);
-
-            // If the block offset is within `shiftBy` bytes of beginning of the block, some bytes we're shifting
-            // will end up in the previous block. Here's an illustrated example:
-            //
-            // shiftBy = 2, blockIndex = 1, blockOffset = 1
-            //                        v---- Cursor is here
-            // Before: [A B C D E] [F G H I J]
-            // After : [A B C D G] [F G H I J]
-            //  Now this is G --^       ^-- And the cursor is here, ready to continue copy the rest.
-            if (blockOffset < shiftBy) {
-                Block previousBlock = blocks.get(blockIndex - 1);
-                int numberOfBytesToShift = Math.min(length, shiftBy) - blockOffset;
-                System.arraycopy(block.data, blockOffset, previousBlock.data, previousBlock.data.length - numberOfBytesToShift, numberOfBytesToShift);
-
-                // Now that we've shifted some bytes, update our position within the buffer.
-                bufferOffset += numberOfBytesToShift;
-                length -= numberOfBytesToShift;
-
-                // If there are no more bytes to shift...
-                if (length == 0) {
-                    // ...lower the `limit` because we've reclaimed some bytes in this block...
-                    block.limit -= numberOfBytesToShift;
-                    // ...and early return.
-                    return;
-                }
-                // Otherwise, use our new buffer offset to recalculate our block-specific position.
-                blockIndex = index(bufferOffset);
-                block = blocks.get(blockIndex);
-                blockOffset = offset(bufferOffset);
-            }
-
-            // At this point, the block offset is at least `shiftBy` bytes away from the beginning of the block, so
-            // we can do a memcpy to shift the rest of the bytes in the block.
-            int numberOfBytesToShift = Math.min(length, block.data.length - blockOffset);
-            System.arraycopy(block.data, blockOffset, block.data, blockOffset - shiftBy, numberOfBytesToShift);
-
-            // Update our counters and see if there are any more bytes to shift.
-            bufferOffset += numberOfBytesToShift;
+        // In this method, "buffer offsets" are absolute indexes into the WriteBuffer and
+        // "block offsets" are indexes that are relative to the beginning of the current Block.
+
+        // The first buffer offset that does not yet contain data.
+        long position = position();
+        // This is the buffer offset of the first byte that we will be shifting backwards.
+        long sourceBufferOffset = position - length;
+        // When we're done, this will be the first offset in the buffer that does not contain data.
+        long writeBufferLimit = position - shiftBy;
+
+        while (length > 0) {
+            // Convert the source buffer offset into a (Block, block offset) pair.
+            int sourceBlockIndex = index(sourceBufferOffset);
+            Block sourceBlock = blocks.get(sourceBlockIndex);
+            int sourceBlockOffset = offset(sourceBufferOffset);
+
+            // Convert the destination buffer offset into a (Block, block offset) pair.
+            // Because buffer offsets are absolute, the `destinationBufferOffset` in each loop iteration is
+            // `shiftBy` positions behind the `sourceBufferOffset`.
+            long destinationBufferOffset = sourceBufferOffset - shiftBy;
+            int destinationBlockIndex = index(destinationBufferOffset);
+            Block destinationBlock = blocks.get(destinationBlockIndex);
+            int destinationBlockOffset = offset(destinationBufferOffset);
+
+            // Determine how many bytes are left in the source and destination blocks following their respective
+            // block offsets.
+            int bytesLeftInSourceBlock = sourceBlock.limit - sourceBlockOffset;
+            int bytesLeftInDestinationBlock = destinationBlock.limit - destinationBlockOffset;
+            // Whichever block has fewer bytes remaining will determine how many bytes we consider to be
+            // available for shifting in this pass.
+            int bytesAvailableToCopy = Math.min(bytesLeftInSourceBlock, bytesLeftInDestinationBlock);
+
+            // If there are more bytes available than we need to finish the shifting operation, take `length` instead.
+            int numberOfBytesToShift = Math.min(length, bytesAvailableToCopy);
+
+            // Copy the bytes from the source to the destination.
+            System.arraycopy(
+                    sourceBlock.data,
+                    sourceBlockOffset,
+                    destinationBlock.data,
+                    destinationBlockOffset,
+                    numberOfBytesToShift
+            );
+
+            // Update our record of how many bytes to shift remain...
             length -= numberOfBytesToShift;
+            // ...and from which point we should resume in the next iteration.
+            sourceBufferOffset += numberOfBytesToShift;
         }
-        if (block != null) {
-            // We've reclaimed some space in this block; lower the `limit` accordingly.
-            block.limit -= shiftBy;
+
+        // At this point, the shifting is complete. However, we have reclaimed some amount of space in the WriteBuffer.
+        // Using the `writeBufferLimit` we calculated at the beginning of the method, find the last Block that still
+        // contains data.
+        int lastBlockIndex = index(writeBufferLimit);
+        Block lastBlock = blocks.get(lastBlockIndex);
+        int lastBlockOffset = offset(writeBufferLimit);
+
+        // Update that Block's limit...
+        lastBlock.limit = lastBlockOffset;
+        // ...and return any empty blocks at the tail of the `blocks` list to the pool.
+        for (int m = blocks.size() - 1; m > lastBlockIndex; m--) {
+            Block emptyBlock = blocks.remove(m);
+            emptyBlock.close();
         }
+
+        // Update the WriteBuffer's member fields to reflect the changes we've made.
+        current = lastBlock;
+        index = lastBlockIndex;
     }
 
     /** Writes an array of bytes to the buffer expanding if necessary, defaulting to the entire array. */


Test patch:
diff --git a/test/com/amazon/ion/impl/bin/WriteBufferTest.java b/test/com/amazon/ion/impl/bin/WriteBufferTest.java
index 24822820..babf113a 100644
--- a/test/com/amazon/ion/impl/bin/WriteBufferTest.java
+++ b/test/com/amazon/ion/impl/bin/WriteBufferTest.java
@@ -959,4 +959,118 @@ public class WriteBufferTest
         buf.shiftBytesLeft(5, 5);
         assertBuffer("012345BCDEF".getBytes());
     }
+
+    @Test
+    public void shiftBytesLeftAcrossBufferBlocksEmptyingLastBlock() throws IOException {
+        assertEquals(11, ALLOCATOR.getBlockSize());
+        // The "B" is the first and only byte in the second block.
+        // Shifting 1 byte left by one empties the last block.
+        buf.writeBytes("0123456789AB".getBytes());
+        buf.shiftBytesLeft(1, 1);
+        assertBuffer("0123456789B".getBytes());
+    }
+
+    @Test
+    public void shiftEntireBlock() throws IOException {
+        assertEquals(11, ALLOCATOR.getBlockSize());
+        // The buffer contains two full blocks
+        buf.writeBytes("0123456789|ABCDEFGHIJ|".getBytes());
+        // We shift an entire block left by the block size
+        buf.shiftBytesLeft(11, 11);
+        assertBuffer("ABCDEFGHIJ|".getBytes());
+    }
+
+    @Test
+    public void shiftBytesLeftByMoreThanTheBlockSize() {
+        assertEquals(11, ALLOCATOR.getBlockSize());
+        // We have 5 blocks' worth of data
+        buf.writeBytes("0123456789|0123456789|0123456789|0123456789|0123456789|".getBytes());
+        // We can shift left amounts greater than the block size
+        buf.shiftBytesLeft(24, 24);
+        assertBuffer("01234569|0123456789|0123456789|".getBytes());
+    }
+
+    @Test
+    public void shiftBytesLeftAcrossBufferBlocksShorteningNextToLastBlock() throws IOException {
+        assertEquals(11, ALLOCATOR.getBlockSize());
+        // The "B" is the first and only byte in the second block.
+        // Shifting 2 bytes left by two empties the last block and shortens the next-to-last block by 1 byte.
+        // Following this operation, the next-to-last block becomes the last block.
+        buf.writeBytes("0123456789AB".getBytes());
+        buf.shiftBytesLeft(2, 2);
+        assertBuffer("01234567AB".getBytes());
+    }
+
+    @Test
+    public void writingAfterLastBlockChanges() throws IOException {
+        assertEquals(11, ALLOCATOR.getBlockSize());
+        // The "B" is the first and only byte in the second block.
+        // Shifting 2 bytes left by two empties the last block and shortens the next-to-last block by 1 byte.
+        // Following this operation, the next-to-last block becomes the last block.
+        buf.writeBytes("0123456789AB".getBytes());
+        buf.shiftBytesLeft(2, 2);
+        assertBuffer("01234567AB".getBytes());
+        // After shifting and changing the last block, we can still append data without issue.
+        buf.writeBytes("CDE".getBytes());
+        assertBuffer("01234567ABCDE".getBytes());
+    }
+
+    @Test
+    public void updateLastBlock() throws IOException {
+        assertEquals(11, ALLOCATOR.getBlockSize());
+        // The "B" is the first and only byte in the second block.
+        // Shifting 2 bytes left by two empties the last block and shortens the next-to-last block by 1 byte.
+        // Following this operation, the next-to-last block becomes the last block.
+        buf.writeBytes("0123456789AB".getBytes());
+        buf.shiftBytesLeft(2, 2);
+        assertBuffer("01234567AB".getBytes());
+        // We write some more data to the buffer. If the last block has been discarded correctly, we can do another
+        // shift operation without getting corrupt data.
+        buf.writeBytes("CDEFGH".getBytes());
+        assertBuffer("01234567ABCDEFGH".getBytes());
+        buf.shiftBytesLeft(4, 2);
+        assertBuffer("01234567ABEFGH".getBytes());
+    }
+
+    @Test
+    public void shiftNBytesLeftByMoreThanNBytes() {
+        assertEquals(11, ALLOCATOR.getBlockSize());
+        buf.writeBytes("0123456789AB".getBytes());
+        // Shift left by more bytes than we're shifting (shiftBy > length)
+        buf.shiftBytesLeft(2, 5);
+        assertBuffer("01234AB".getBytes());
+    }
+
+    @Test
+    public void shiftLeftToBeginning() {
+        assertEquals(11, ALLOCATOR.getBlockSize());
+        buf.writeBytes("01234567AB".getBytes());
+        buf.shiftBytesLeft(2, 8);
+        assertBuffer("AB".getBytes());
+    }
+
+    @Test
+    public void shiftLeftToBeginningAcrossBlocks() {
+        assertEquals(11, ALLOCATOR.getBlockSize());
+        buf.writeBytes("0123456789AB".getBytes());
+        buf.shiftBytesLeft(2, 10);
+        assertBuffer("AB".getBytes());
+    }
+
+    @Test
+    public void shiftBytesLeftAcrossBufferBlocksToStartThenWritePastBlockBoundary() throws IOException {
+        assertEquals(11, ALLOCATOR.getBlockSize());
+        // The "B" is the first and only byte in the second block.
+        // Shifting 2 bytes left by 10 empties the last block and positions the shifted bytes at the start of the first
+        // block.
+        buf.writeBytes("0123456789AB".getBytes());
+        buf.shiftBytesLeft(2, 10);
+        assertBuffer("AB".getBytes());
+        // Write enough bytes to expand back into the next block.
+        buf.writeBytes("CDE0123456789AB".getBytes());
+        assertBuffer("ABCDE0123456789AB".getBytes());
+        // Shift one byte left past the block boundary. This should once again empty the last block.
+        buf.shiftBytesLeft(1, 6);
+        assertBuffer("ABCDE01234B".getBytes());
+    }
 }