chore: update ChunkSegmenter to optionally allow a limit on the number of bytes it should consume (#3279)

* chore: update ChunkSegmenter to optionally allow a limit on the number of bytes it should consume

Update BidiAppendableUnbufferedWritableByteChannel to only attempt to consume as many bytes as are available according to the stream -- this prevents over packing of segments we will for sure never be able to use.

* chore: fix blockSize reset logic

Author: BenWhitehead

google-cloud-storage/src/main/java/com/google/cloud/storage/BidiAppendableUnbufferedWritableByteChannel.java

@@ -136,7 +136,8 @@ private long internalWrite(ByteBuffer[] srcs, int srcsOffset, int srcsLength) th
     RewindableContent rewindableContent = RewindableContent.of(srcs, srcsOffset, srcsLength);
     long totalBufferRemaining = rewindableContent.getLength();
 
-    ChunkSegment[] data = chunkSegmenter.segmentBuffers(srcs, srcsOffset, srcsLength, true);
+    ChunkSegment[] data =
+        chunkSegmenter.segmentBuffers(srcs, srcsOffset, srcsLength, true, availableCapacity);
     if (data.length == 0) {
       return 0;
     }

google-cloud-storage/src/main/java/com/google/cloud/storage/ChunkSegmenter.java

@@ -18,8 +18,10 @@
 
 import com.google.cloud.storage.Crc32cValue.Crc32cLengthKnown;
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.MoreObjects;
 import com.google.common.base.Preconditions;
 import com.google.common.math.IntMath;
+import com.google.common.primitives.Ints;
 import com.google.protobuf.ByteString;
 import java.math.RoundingMode;
 import java.nio.ByteBuffer;
@@ -97,66 +99,96 @@ ChunkSegment[] segmentBuffers(
     // turn this into a single branch, rather than multiple that would need to be checked each
     // element of the iteration
     if (allowUnalignedBlocks) {
-      return segmentWithUnaligned(bbs, offset, length);
+      return segmentWithUnaligned(bbs, offset, length, Long.MAX_VALUE);
     } else {
-      return segmentWithoutUnaligned(bbs, offset, length);
+      return segmentWithoutUnaligned(bbs, offset, length, Long.MAX_VALUE);
     }
   }
 
-  private ChunkSegment[] segmentWithUnaligned(ByteBuffer[] bbs, int offset, int length) {
+  ChunkSegment[] segmentBuffers(
+      ByteBuffer[] bbs,
+      int offset,
+      int length,
+      boolean allowUnalignedBlocks,
+      long maxBytesToConsume) {
+    // turn this into a single branch, rather than multiple that would need to be checked each
+    // element of the iteration
+    if (allowUnalignedBlocks) {
+      return segmentWithUnaligned(bbs, offset, length, maxBytesToConsume);
+    } else {
+      long misaligned = maxBytesToConsume % blockSize;
+      long alignedMaxBytesToConsume = maxBytesToConsume - misaligned;
+      return segmentWithoutUnaligned(bbs, offset, length, alignedMaxBytesToConsume);
+    }
+  }
+
+  private ChunkSegment[] segmentWithUnaligned(
+      ByteBuffer[] bbs, int offset, int length, long maxBytesToConsume) {
     Deque<ChunkSegment> data = new ArrayDeque<>();
 
+    long consumed = 0;
     for (int i = offset; i < length; i++) {
       ByteBuffer buffer = bbs[i];
       int remaining;
-      while ((remaining = buffer.remaining()) > 0) {
-        consumeBytes(data, remaining, buffer);
+      while ((remaining = buffer.remaining()) > 0 && consumed < maxBytesToConsume) {
+        long remainingConsumable = maxBytesToConsume - consumed;
+        int toConsume = remaining;
+        if (remainingConsumable < remaining) {
+          toConsume = Math.toIntExact(remainingConsumable);
+        }
+        long consumeBytes = consumeBytes(data, toConsume, buffer);
+        consumed += consumeBytes;
       }
     }
 
     return data.toArray(new ChunkSegment[0]);
   }
 
-  private ChunkSegment[] segmentWithoutUnaligned(ByteBuffer[] bbs, int offset, int length) {
+  private ChunkSegment[] segmentWithoutUnaligned(
+      ByteBuffer[] bbs, int offset, int length, long maxBytesToConsume) {
     Deque<ChunkSegment> data = new ArrayDeque<>();
 
-    final long totalRemaining = Buffers.totalRemaining(bbs, offset, length);
+    long buffersTotalRemaining = Buffers.totalRemaining(bbs, offset, length);
+    final long totalRemaining = Math.min(maxBytesToConsume, buffersTotalRemaining);
     long consumedSoFar = 0;
 
     int currentBlockPending = blockSize;
 
+    outerloop:
     for (int i = offset; i < length; i++) {
       ByteBuffer buffer = bbs[i];
       int remaining;
       while ((remaining = buffer.remaining()) > 0) {
         long overallRemaining = totalRemaining - consumedSoFar;
         if (overallRemaining < blockSize && currentBlockPending == blockSize) {
-          break;
+          break outerloop;
         }
 
         int numBytesConsumable;
-        if (remaining >= blockSize) {
+        if (remaining >= blockSize && currentBlockPending == blockSize) {
           int blockCount = IntMath.divide(remaining, blockSize, RoundingMode.DOWN);
           numBytesConsumable = blockCount * blockSize;
-        } else if (currentBlockPending < blockSize) {
-          numBytesConsumable = currentBlockPending;
-          currentBlockPending = blockSize;
         } else {
-          numBytesConsumable = remaining;
-          currentBlockPending = currentBlockPending - remaining;
+          numBytesConsumable = Math.min(remaining, currentBlockPending);
         }
         if (numBytesConsumable <= 0) {
-          continue;
+          break outerloop;
         }
 
-        consumedSoFar += consumeBytes(data, numBytesConsumable, buffer);
+        int consumed = consumeBytes(data, numBytesConsumable, buffer);
+        int currentBlockPendingLessConsumed = currentBlockPending - consumed;
+        currentBlockPending = currentBlockPendingLessConsumed % blockSize;
+        if (currentBlockPending == 0) {
+          currentBlockPending = blockSize;
+        }
+        consumedSoFar += consumed;
       }
     }
 
     return data.toArray(new ChunkSegment[0]);
   }
 
-  private long consumeBytes(Deque<ChunkSegment> data, int numBytesConsumable, ByteBuffer buffer) {
+  private int consumeBytes(Deque<ChunkSegment> data, int numBytesConsumable, ByteBuffer buffer) {
     // either no chunk or most recent chunk is full, start a new one
     ChunkSegment peekLast = data.peekLast();
     if (peekLast == null || peekLast.b.size() == maxSegmentSize) {
@@ -167,7 +199,8 @@ private long consumeBytes(Deque<ChunkSegment> data, int numBytesConsumable, Byte
     } else {
       ChunkSegment chunkSoFar = data.pollLast();
       //noinspection ConstantConditions -- covered by peekLast check above
-      int limit = Math.min(numBytesConsumable, maxSegmentSize - chunkSoFar.b.size());
+      int limit =
+          Ints.min(buffer.remaining(), numBytesConsumable, maxSegmentSize - chunkSoFar.b.size());
       ChunkSegment datum = newSegment(buffer, limit);
       ChunkSegment plus = chunkSoFar.concat(datum);
       data.addLast(plus);
@@ -218,5 +251,14 @@ public Crc32cLengthKnown getCrc32c() {
     public boolean isOnlyFullBlocks() {
       return onlyFullBlocks;
     }
+
+    @Override
+    public String toString() {
+      return MoreObjects.toStringHelper(this)
+          .add("crc32c", crc32c)
+          .add("onlyFullBlocks", onlyFullBlocks)
+          .add("b", b)
+          .toString();
+    }
   }
 }

google-cloud-storage/src/test/java/com/google/cloud/storage/ChunkSegmenterTest.java

@@ -21,6 +21,7 @@
 
 import com.google.cloud.storage.ChunkSegmenter.ChunkSegment;
 import com.google.cloud.storage.Crc32cValue.Crc32cLengthKnown;
+import com.google.cloud.storage.it.ChecksummedTestContent;
 import com.google.common.collect.ImmutableList;
 import com.google.common.hash.HashCode;
 import com.google.common.hash.Hashing;
@@ -172,6 +173,78 @@ void allowUnalignedBlocks_false_3() throws Exception {
         () -> assertThat(actual).isEqualTo(expected));
   }
 
+  @Example
+  void maxBytesToConsume_unaligned() throws Exception {
+
+    ChecksummedTestContent ctc = ChecksummedTestContent.gen(64);
+
+    ChunkSegmenter segmenter = new ChunkSegmenter(Hasher.noop(), ByteStringStrategy.noCopy(), 6, 3);
+
+    List<ChecksummedTestContent> chunks = ctc.chunkup(4);
+    ByteBuffer[] buffers =
+        chunks.stream().map(ChecksummedTestContent::asByteBuffer).toArray(ByteBuffer[]::new);
+    buffers[1].position(1);
+
+    ChecksummedTestContent slice = ctc.slice(5, 37);
+    List<ByteString> expected =
+        slice.chunkup(6).stream()
+            .map(ChecksummedTestContent::asByteBuffer)
+            .map(ByteStringStrategy.noCopy())
+            .collect(Collectors.toList());
+
+    ChunkSegment[] segments = segmenter.segmentBuffers(buffers, 1, buffers.length - 2, true, 37);
+    List<ByteString> actual =
+        Arrays.stream(segments).map(ChunkSegment::getB).collect(Collectors.toList());
+    assertThat(actual).isEqualTo(expected);
+  }
+
+  @Example
+  void maxBytesToConsume_aligned() throws Exception {
+
+    ChecksummedTestContent ctc = ChecksummedTestContent.gen(64);
+
+    ChunkSegmenter segmenter = new ChunkSegmenter(Hasher.noop(), ByteStringStrategy.noCopy(), 6, 3);
+
+    List<ChecksummedTestContent> chunks = ctc.chunkup(4);
+    ByteBuffer[] buffers =
+        chunks.stream().map(ChecksummedTestContent::asByteBuffer).toArray(ByteBuffer[]::new);
+    buffers[1].position(1);
+
+    ChecksummedTestContent slice = ctc.slice(5, 36);
+    List<ByteString> expected =
+        slice.chunkup(6).stream()
+            .map(ChecksummedTestContent::asByteBuffer)
+            .map(ByteStringStrategy.noCopy())
+            .collect(Collectors.toList());
+
+    ChunkSegment[] segments = segmenter.segmentBuffers(buffers, 1, buffers.length - 2, false, 37);
+    List<ByteString> actual =
+        Arrays.stream(segments).map(ChunkSegment::getB).collect(Collectors.toList());
+    assertThat(actual).isEqualTo(expected);
+  }
+
+  @Example
+  void alignedConsumeForLargeBuffersOnlyConsumesAligned() throws Exception {
+
+    ChecksummedTestContent ctc = ChecksummedTestContent.gen(2048 + 13);
+
+    ChunkSegmenter segmenter =
+        new ChunkSegmenter(Hasher.noop(), ByteStringStrategy.noCopy(), 2048, 256);
+
+    ChecksummedTestContent slice = ctc.slice(0, 2048);
+    List<ByteString> expected =
+        slice.chunkup(2048).stream()
+            .map(ChecksummedTestContent::asByteBuffer)
+            .map(ByteStringStrategy.noCopy())
+            .collect(Collectors.toList());
+
+    ByteBuffer buf = ctc.asByteBuffer();
+    ChunkSegment[] segments = segmenter.segmentBuffers(new ByteBuffer[] {buf}, 0, 1, false);
+    List<ByteString> actual =
+        Arrays.stream(segments).map(ChunkSegment::getB).collect(Collectors.toList());
+    assertThat(actual).isEqualTo(expected);
+  }
+
   @Provide("TestData")
   static Arbitrary<TestData> arbitraryTestData() {
     return Arbitraries.lazyOf(

google-cloud-storage/src/test/java/com/google/cloud/storage/ITAppendableUploadTest.java

@@ -263,7 +263,8 @@ public void takeoverJustToFinalizeWorks() throws Exception {
   }
 
   private void checkTestbenchIssue733() {
-    if (p.uploadConfig.getCloseAction() == CloseAction.FINALIZE_WHEN_CLOSING) {
+    if (backend == Backend.TEST_BENCH
+        && p.uploadConfig.getCloseAction() == CloseAction.FINALIZE_WHEN_CLOSING) {
       int estimatedMessageCount = 0;
       FlushPolicy flushPolicy = p.uploadConfig.getFlushPolicy();
       if (flushPolicy instanceof MinFlushSizeFlushPolicy) {