Repository: googleapis/java-storage
Pull request #1467 (2022-06-22 22:12:37)

chore: refactor Buffered Channels such that buffer allocation is now lazy
Eagar allocation of large buffers can cause significant memory pressure for users
who are processing many objects in parallel or processing small objects. Since our
Buffered Channels are present to smooth out rpc sizes which can have high latencies
we do not want to unnecessarily allocate a buffer.

Buffered Read & Write Channel already minimize unnecessary copies into our buffer,
coupled with this change, if a user is reading/writing buffers the same size as the
configured buffer size we will never actually allocate a buffer thereby saving on
unnecessary memory consumption.



Patch:
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/BufferHandle.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/BufferHandle.java
new file mode 100644
index 00000000..d230e9d6
--- /dev/null
+++ b/google-cloud-storage/src/main/java/com/google/cloud/storage/BufferHandle.java
@@ -0,0 +1,146 @@
+/*
+ * Copyright 2022 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.storage;
+
+import com.google.common.annotations.VisibleForTesting;
+import java.nio.ByteBuffer;
+import java.util.function.Function;
+import java.util.function.Supplier;
+
+/**
+ * Sometimes, we need a handle to create a buffer but do not want to unnecessarily allocate it. This
+ * class can be though of as an enriched {@code Supplier<ByteBuffer>} that still implements the
+ * common stateful methods of {@link ByteBuffer} without allocating the ByteBuffer until necessary.
+ *
+ * <p>{@link ByteBuffer} is a sealed class hierarchy, meaning we can't simply extend it to provide
+ * laziness without this new class.
+ */
+abstract class BufferHandle implements Supplier<ByteBuffer> {
+
+  @VisibleForTesting
+  BufferHandle() {}
+
+  abstract int remaining();
+
+  abstract int capacity();
+
+  abstract int position();
+
+  static BufferHandle allocateAligned(int alignmentMultiple, int size) {
+    int actualSize = alignSize(alignmentMultiple, size);
+    return allocate(actualSize);
+  }
+
+  static BufferHandle allocate(int capacity) {
+    return new LazyBufferHandle(capacity, ByteBuffer::allocate);
+  }
+
+  static BufferHandle handleOf(ByteBuffer buf) {
+    return new EagerBufferHandle(buf);
+  }
+
+  /**
+   * Give {@code size} "snap" it to the next {@code alignmentMultiple} that is >= {@code size}.
+   *
+   * <p>i.e. Given 344k size, 256k alignmentMultiple expect 512k
+   */
+  private static int alignSize(int alignmentMultiple, int size) {
+    int actualSize = size;
+    if (size < alignmentMultiple) {
+      actualSize = alignmentMultiple;
+    } else if (size % alignmentMultiple != 0) {
+      // TODO: this mod will cause two divisions to happen
+      //   * try and measure how expensive two divisions is compared to one
+      //   * also measure the case where size is a multiple, and how much the
+      //     following calculation costs
+
+      // add almost another full alignmentMultiple to the size
+      // then integer divide it before multiplying it by the alignmentMultiple
+      actualSize = (size + alignmentMultiple - 1) / alignmentMultiple * alignmentMultiple;
+    } // else size is already aligned
+    return actualSize;
+  }
+
+  static final class LazyBufferHandle extends BufferHandle {
+
+    private final int capacity;
+    private final Function<Integer, ByteBuffer> factory;
+
+    private volatile ByteBuffer buf;
+
+    @VisibleForTesting
+    LazyBufferHandle(int capacity, Function<Integer, ByteBuffer> factory) {
+      this.capacity = capacity;
+      this.factory = factory;
+    }
+
+    @Override
+    int remaining() {
+      return buf == null ? capacity() : buf.remaining();
+    }
+
+    @Override
+    int capacity() {
+      return buf == null ? capacity : buf.capacity();
+    }
+
+    @Override
+    int position() {
+      return buf == null ? 0 : buf.position();
+    }
+
+    @Override
+    public ByteBuffer get() {
+      if (buf == null) {
+        synchronized (this) {
+          if (buf == null) {
+            buf = factory.apply(capacity);
+          }
+        }
+      }
+      return buf;
+    }
+  }
+
+  static final class EagerBufferHandle extends BufferHandle {
+    private final ByteBuffer buf;
+
+    private EagerBufferHandle(ByteBuffer buf) {
+      this.buf = buf;
+    }
+
+    @Override
+    int remaining() {
+      return buf.remaining();
+    }
+
+    @Override
+    int capacity() {
+      return buf.capacity();
+    }
+
+    @Override
+    int position() {
+      return buf.position();
+    }
+
+    @Override
+    public ByteBuffer get() {
+      return buf;
+    }
+  }
+}
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/DefaultBufferedReadableByteChannel.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/DefaultBufferedReadableByteChannel.java
index f396d8e5..45359a22 100644
--- a/google-cloud-storage/src/main/java/com/google/cloud/storage/DefaultBufferedReadableByteChannel.java
+++ b/google-cloud-storage/src/main/java/com/google/cloud/storage/DefaultBufferedReadableByteChannel.java
@@ -24,15 +24,15 @@ import java.nio.channels.ClosedChannelException;
 
 final class DefaultBufferedReadableByteChannel implements BufferedReadableByteChannel {
 
-  private final ByteBuffer buffer;
+  private final BufferHandle handle;
 
   private final UnbufferedReadableByteChannel channel;
 
   private boolean flipped = false;
   private boolean retEOF = false;
 
-  DefaultBufferedReadableByteChannel(ByteBuffer buffer, UnbufferedReadableByteChannel channel) {
-    this.buffer = buffer;
+  DefaultBufferedReadableByteChannel(BufferHandle handle, UnbufferedReadableByteChannel channel) {
+    this.handle = handle;
     this.channel = channel;
   }
 
@@ -48,13 +48,14 @@ final class DefaultBufferedReadableByteChannel implements BufferedReadableByteCh
     int bytesConsumed = 0;
 
     while (dst.hasRemaining()) {
-      int bufferRemaining = buffer.remaining();
+      int bufferRemaining = handle.remaining();
 
       int dstRemaining = dst.remaining();
       int dstPosition = dst.position();
 
       final int tmpBytesCopied;
       if (enqueuedBytes()) {
+        ByteBuffer buffer = handle.get();
         if (!flipped) {
           buffer.flip();
           flipped = true;
@@ -96,6 +97,7 @@ final class DefaultBufferedReadableByteChannel implements BufferedReadableByteCh
           // create a slice of our buffer such that
           // dst.remaning() + bufSlice.remaning() == buffer.capacity
 
+          ByteBuffer buffer = handle.get();
           ByteBuffer slice = buffer.slice();
           int sliceCapacity = buffer.capacity() - dstRemaining;
           Buffers.limit(slice, sliceCapacity);
@@ -140,6 +142,6 @@ final class DefaultBufferedReadableByteChannel implements BufferedReadableByteCh
   }
 
   private boolean enqueuedBytes() {
-    return buffer.position() > 0;
+    return handle.position() > 0;
   }
 }
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/DefaultBufferedWritableByteChannel.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/DefaultBufferedWritableByteChannel.java
index 73368fc1..5f1834bb 100644
--- a/google-cloud-storage/src/main/java/com/google/cloud/storage/DefaultBufferedWritableByteChannel.java
+++ b/google-cloud-storage/src/main/java/com/google/cloud/storage/DefaultBufferedWritableByteChannel.java
@@ -55,12 +55,12 @@ import java.nio.channels.ClosedChannelException;
  */
 final class DefaultBufferedWritableByteChannel implements BufferedWritableByteChannel {
 
-  private final ByteBuffer buffer;
+  private final BufferHandle handle;
 
   private final UnbufferedWritableByteChannel channel;
 
-  DefaultBufferedWritableByteChannel(ByteBuffer buffer, UnbufferedWritableByteChannel channel) {
-    this.buffer = buffer;
+  DefaultBufferedWritableByteChannel(BufferHandle handle, UnbufferedWritableByteChannel channel) {
+    this.handle = handle;
     this.channel = channel;
   }
 
@@ -76,14 +76,14 @@ final class DefaultBufferedWritableByteChannel implements BufferedWritableByteCh
       int srcRemaining = src.remaining();
       int srcPosition = src.position();
 
-      int capacity = buffer.capacity();
-      int bufferRemaining = buffer.remaining();
+      int capacity = handle.capacity();
+      int bufferRemaining = handle.remaining();
       int bufferPending = capacity - bufferRemaining;
 
-      boolean enqueuedBytes = buffer.hasRemaining();
-      if (enqueuedBytes && srcRemaining < bufferRemaining) {
+      boolean enqueuedBytes = enqueuedBytes();
+      if (srcRemaining < bufferRemaining) {
         // srcRemaining is smaller than the remaining space in our buffer, enqueue it in full
-        buffer.put(src);
+        handle.get().put(src);
         bytesConsumed += srcRemaining;
         break;
       } else if (enqueuedBytes) {
@@ -105,6 +105,7 @@ final class DefaultBufferedWritableByteChannel implements BufferedWritableByteCh
           buf = slice;
         }
 
+        ByteBuffer buffer = handle.get();
         Buffers.flip(buffer);
         ByteBuffer[] srcs = {buffer, buf};
         long write = channel.write(srcs);
@@ -131,13 +132,13 @@ final class DefaultBufferedWritableByteChannel implements BufferedWritableByteCh
           }
         }
       } else {
-        // no enqueued data, see if we can simply write the provided src or a slice of it
-        // since our buffer is empty
+        // no enqueued data and src is at least as large as our buffer, see if we can simply write
+        // the provided src or a slice of it since our buffer is empty
         if (bufferRemaining == srcRemaining) {
           // the capacity of buffer and the bytes remaining in src are the same, directly
           // write src
           bytesConsumed += channel.write(src);
-        } else if (bufferRemaining <= srcRemaining) {
+        } else {
           // the src provided is larger than our buffer. rather than copying into the buffer, simply
           // write a slice
           ByteBuffer slice = src.slice();
@@ -145,10 +146,6 @@ final class DefaultBufferedWritableByteChannel implements BufferedWritableByteCh
           int write = channel.write(slice);
           Buffers.position(src, srcPosition + write);
           bytesConsumed += write;
-        } else {
-          // srcRemaining is smaller than the remaining space in our buffer, enqueue it in full
-          buffer.put(src);
-          bytesConsumed += srcRemaining;
         }
       }
     }
@@ -169,7 +166,8 @@ final class DefaultBufferedWritableByteChannel implements BufferedWritableByteCh
 
   @Override
   public void flush() throws IOException {
-    if (buffer.hasRemaining()) {
+    if (enqueuedBytes()) {
+      ByteBuffer buffer = handle.get();
       Buffers.flip(buffer);
       channel.write(buffer);
       if (buffer.hasRemaining()) {
@@ -179,4 +177,8 @@ final class DefaultBufferedWritableByteChannel implements BufferedWritableByteCh
       }
     }
   }
+
+  private boolean enqueuedBytes() {
+    return handle.position() > 0;
+  }
 }
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/GapicDownloadSessionBuilder.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/GapicDownloadSessionBuilder.java
index e5574703..25e592e2 100644
--- a/google-cloud-storage/src/main/java/com/google/cloud/storage/GapicDownloadSessionBuilder.java
+++ b/google-cloud-storage/src/main/java/com/google/cloud/storage/GapicDownloadSessionBuilder.java
@@ -62,8 +62,14 @@ final class GapicDownloadSessionBuilder {
       return this;
     }
 
+    public BufferedReadableByteChannelSessionBuilder buffered(int capacity) {
+      return new BufferedReadableByteChannelSessionBuilder(
+          BufferHandle.allocate(capacity), getF(read, hasher));
+    }
+
     public BufferedReadableByteChannelSessionBuilder buffered(ByteBuffer buffer) {
-      return new BufferedReadableByteChannelSessionBuilder(buffer, getF(read, hasher));
+      return new BufferedReadableByteChannelSessionBuilder(
+          BufferHandle.handleOf(buffer), getF(read, hasher));
     }
 
     public UnbufferedReadableByteChannelSessionBuilder unbuffered() {
@@ -82,7 +88,7 @@ final class GapicDownloadSessionBuilder {
       private Object obj;
 
       private BufferedReadableByteChannelSessionBuilder(
-          ByteBuffer buffer,
+          BufferHandle buffer,
           BiFunction<Object, SettableApiFuture<Object>, UnbufferedReadableByteChannel> f) {
         this.f = f.andThen(c -> new DefaultBufferedReadableByteChannel(buffer, c));
       }
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/GapicResumableUploadSessionBuilder.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/GapicResumableUploadSessionBuilder.java
index 6eb47655..62ff18db 100644
--- a/google-cloud-storage/src/main/java/com/google/cloud/storage/GapicResumableUploadSessionBuilder.java
+++ b/google-cloud-storage/src/main/java/com/google/cloud/storage/GapicResumableUploadSessionBuilder.java
@@ -72,7 +72,14 @@ final class GapicResumableUploadSessionBuilder {
 
     public BufferedWritableByteChannelSessionBuilder buffered(ByteBuffer buffer) {
       return new BufferedWritableByteChannelSessionBuilder(
-          buffer, getF(write, byteStringStrategy, hasher));
+          BufferHandle.handleOf(buffer), getF(write, byteStringStrategy, hasher));
+    }
+
+    public BufferedWritableByteChannelSessionBuilder bufferedAndAligned(
+        int alignmentMultiple, int capacity) {
+      return new BufferedWritableByteChannelSessionBuilder(
+          BufferHandle.allocateAligned(alignmentMultiple, capacity),
+          getF(write, byteStringStrategy, hasher));
     }
 
     public UnbufferedWritableByteChannelSessionBuilder unbuffered() {
@@ -99,7 +106,7 @@ final class GapicResumableUploadSessionBuilder {
       private ApiFuture<ResumableWrite> uploadIdFuture;
 
       private BufferedWritableByteChannelSessionBuilder(
-          ByteBuffer buffer,
+          BufferHandle buffer,
           BiFunction<
                   ResumableWrite,
                   SettableApiFuture<WriteObjectResponse>,


Test patch:
diff --git a/google-cloud-storage/src/test/java/com/google/cloud/storage/BufferHandleTest.java b/google-cloud-storage/src/test/java/com/google/cloud/storage/BufferHandleTest.java
new file mode 100644
index 00000000..d88b56b2
--- /dev/null
+++ b/google-cloud-storage/src/test/java/com/google/cloud/storage/BufferHandleTest.java
@@ -0,0 +1,114 @@
+/*
+ * Copyright 2022 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.storage;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import com.google.cloud.storage.BufferHandle.LazyBufferHandle;
+import java.nio.ByteBuffer;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.atomic.AtomicBoolean;
+import org.junit.Assert;
+import org.junit.Test;
+
+public final class BufferHandleTest {
+
+  @Test
+  public void lazyBufferHandle_methodsBehaveTheSameAsAnEmptyByteBuffer() {
+    int capacity = 10;
+    ByteBuffer baseline = ByteBuffer.allocate(capacity);
+    LazyBufferHandle handle =
+        new LazyBufferHandle(
+            capacity,
+            i -> {
+              Assert.fail("should not be called");
+              return null;
+            });
+
+    assertThat(handle.remaining()).isEqualTo(baseline.remaining());
+    assertThat(handle.position()).isEqualTo(baseline.position());
+    assertThat(handle.capacity()).isEqualTo(baseline.capacity());
+  }
+
+  @Test
+  public void lazyBufferHandle_afterAllocationOnlyBackingIsReferenced() {
+    int capacity = 10;
+    ByteBuffer baseline = ByteBuffer.allocate(capacity);
+    AtomicBoolean alloc = new AtomicBoolean(false);
+    LazyBufferHandle handle =
+        new LazyBufferHandle(
+            capacity,
+            i -> {
+              alloc.compareAndSet(false, true);
+              return ByteBuffer.allocate(capacity);
+            });
+
+    assertThat(handle.remaining()).isEqualTo(baseline.remaining());
+    assertThat(handle.position()).isEqualTo(baseline.position());
+    assertThat(handle.capacity()).isEqualTo(baseline.capacity());
+
+    byte[] bytes = new byte[] {(byte) 'a', (byte) 'b'};
+    handle.get().put(bytes);
+    assertThat(alloc.get()).isTrue();
+    assertThat(handle.remaining()).isEqualTo(8);
+    assertThat(handle.position()).isEqualTo(2);
+    assertThat(handle.capacity()).isEqualTo(capacity);
+  }
+
+  @Test
+  public void lazyBufferHandle_initIsThreadSafe() throws ExecutionException, InterruptedException {
+    int capacity = 10;
+    ExecutorService exec = Executors.newFixedThreadPool(2);
+    AtomicBoolean alloc = new AtomicBoolean(false);
+    LazyBufferHandle handle =
+        new LazyBufferHandle(
+            capacity,
+            i -> {
+              alloc.compareAndSet(false, true);
+              return ByteBuffer.allocate(capacity);
+            });
+
+    Future<ByteBuffer> f1 = exec.submit(handle::get);
+    Future<ByteBuffer> f2 = exec.submit(handle::get);
+
+    assertThat(f1.get()).isSameInstanceAs(f2.get());
+
+    assertThat(handle.get().capacity()).isEqualTo(capacity);
+  }
+
+  @Test
+  public void eagerBufferHandle_methodsBehaveTheSameAsAnEmptyByteBuffer() {
+    int capacity = 10;
+    ByteBuffer baseline = ByteBuffer.allocate(capacity);
+    BufferHandle handle = BufferHandle.handleOf(baseline);
+
+    assertThat(handle.remaining()).isEqualTo(baseline.remaining());
+    assertThat(handle.position()).isEqualTo(baseline.position());
+    assertThat(handle.capacity()).isEqualTo(baseline.capacity());
+
+    byte[] bytes = new byte[] {(byte) 'a', (byte) 'b'};
+    baseline.put(bytes);
+    assertThat(handle.remaining()).isEqualTo(8);
+    assertThat(handle.position()).isEqualTo(2);
+    assertThat(handle.capacity()).isEqualTo(capacity);
+
+    assertThat(handle.get()).isSameInstanceAs(baseline);
+  }
+}
diff --git a/google-cloud-storage/src/test/java/com/google/cloud/storage/DefaultBufferedReadableByteChannelTest.java b/google-cloud-storage/src/test/java/com/google/cloud/storage/DefaultBufferedReadableByteChannelTest.java
index 094e8ca4..18c9799f 100644
--- a/google-cloud-storage/src/test/java/com/google/cloud/storage/DefaultBufferedReadableByteChannelTest.java
+++ b/google-cloud-storage/src/test/java/com/google/cloud/storage/DefaultBufferedReadableByteChannelTest.java
@@ -54,7 +54,7 @@ public final class DefaultBufferedReadableByteChannelTest {
     byte[] bytes = DataGenerator.base64Characters().genBytes(61);
 
     ByteBuffer buf = ByteBuffer.allocate(16);
-    ByteBuffer buffer = ByteBuffer.allocate(3);
+    BufferHandle buffer = BufferHandle.allocate(3);
     try (ByteArrayInputStream bais = new ByteArrayInputStream(bytes);
         UnbufferedReadableByteChannelAdapter adapter =
             new UnbufferedReadableByteChannelAdapter(Channels.newChannel(bais));
@@ -83,7 +83,7 @@ public final class DefaultBufferedReadableByteChannelTest {
     byte[] bytes = readOps.bytes;
 
     ByteBuffer buf = ByteBuffer.allocate(readOps.readSize);
-    ByteBuffer buffer = ByteBuffer.allocate(readOps.bufferSize);
+    BufferHandle buffer = BufferHandle.allocate(readOps.bufferSize);
     try (ByteArrayInputStream bais = new ByteArrayInputStream(bytes);
         UnbufferedReadableByteChannelAdapter adapter =
             new UnbufferedReadableByteChannelAdapter(Channels.newChannel(bais));
diff --git a/google-cloud-storage/src/test/java/com/google/cloud/storage/DefaultBufferedWritableByteChannelTest.java b/google-cloud-storage/src/test/java/com/google/cloud/storage/DefaultBufferedWritableByteChannelTest.java
index 0f52a50a..23d5b984 100644
--- a/google-cloud-storage/src/test/java/com/google/cloud/storage/DefaultBufferedWritableByteChannelTest.java
+++ b/google-cloud-storage/src/test/java/com/google/cloud/storage/DefaultBufferedWritableByteChannelTest.java
@@ -56,10 +56,11 @@ public final class DefaultBufferedWritableByteChannelTest {
   @Property
   void bufferingEagerlyFlushesWhenFull(@ForAll("WriteOps") WriteOps writeOps) throws IOException {
     ByteBuffer buffer = ByteBuffer.allocate(writeOps.bufferSize);
+    AuditingBufferHandle handle = new AuditingBufferHandle(BufferHandle.handleOf(buffer));
     try (ByteArrayOutputStream baos = new ByteArrayOutputStream();
         CountingWritableByteChannelAdapter adapter =
             new CountingWritableByteChannelAdapter(Channels.newChannel(baos));
-        BufferedWritableByteChannel c = new DefaultBufferedWritableByteChannel(buffer, adapter)) {
+        BufferedWritableByteChannel c = new DefaultBufferedWritableByteChannel(handle, adapter)) {
 
       List<Integer> actualWriteSizes = new ArrayList<>();
 
@@ -149,8 +150,6 @@ public final class DefaultBufferedWritableByteChannelTest {
    */
   @Example
   void partialFlushOfEnqueuedBytesFlushesMultipleTimes() throws IOException {
-    ByteBuffer buffer = ByteBuffer.allocate(10);
-
     ByteBuffer data1 = DataGenerator.base64Characters().genByteBuffer(5);
     ByteBuffer data2 = DataGenerator.base64Characters().genByteBuffer(9);
     ByteBuffer data3 = DataGenerator.base64Characters().genByteBuffer(3);
@@ -163,10 +162,11 @@ public final class DefaultBufferedWritableByteChannelTest {
         buffers.stream().reduce(ByteBuffer.allocate(allDataSize), ByteBuffer::put).array();
     buffers.forEach(b -> b.position(0));
 
+    AuditingBufferHandle handle = new AuditingBufferHandle(BufferHandle.allocate(10));
     try (ByteArrayOutputStream baos = new ByteArrayOutputStream();
         CountingWritableByteChannelAdapter adapter =
             new CountingWritableByteChannelAdapter(Channels.newChannel(baos));
-        BufferedWritableByteChannel c = new DefaultBufferedWritableByteChannel(buffer, adapter)) {
+        BufferedWritableByteChannel c = new DefaultBufferedWritableByteChannel(handle, adapter)) {
 
       c.write(data1); // write 5 bytes, which should enqueue in full
       // before the next write, limit the number of bytes the underlying channel will consume to 2.
@@ -217,8 +217,6 @@ public final class DefaultBufferedWritableByteChannelTest {
    */
   @Example
   void manualFlushingIsAccurate() throws IOException {
-    ByteBuffer buffer = ByteBuffer.allocate(5);
-
     ByteBuffer data1 = DataGenerator.base64Characters().genByteBuffer(3);
     ByteBuffer data2 = DataGenerator.base64Characters().genByteBuffer(3);
     ByteBuffer data3 = DataGenerator.base64Characters().genByteBuffer(3);
@@ -231,10 +229,11 @@ public final class DefaultBufferedWritableByteChannelTest {
         buffers.stream().reduce(ByteBuffer.allocate(allDataSize), ByteBuffer::put).array();
     buffers.forEach(b -> b.position(0));
 
+    AuditingBufferHandle handle = new AuditingBufferHandle(BufferHandle.allocate(5));
     try (ByteArrayOutputStream baos = new ByteArrayOutputStream();
         CountingWritableByteChannelAdapter adapter =
             new CountingWritableByteChannelAdapter(Channels.newChannel(baos));
-        BufferedWritableByteChannel c = new DefaultBufferedWritableByteChannel(buffer, adapter)) {
+        BufferedWritableByteChannel c = new DefaultBufferedWritableByteChannel(handle, adapter)) {
 
       c.write(data1); // write 3 bytes, which should enqueue in full
       c.flush(); // flush all enqueued bytes
@@ -344,6 +343,43 @@ public final class DefaultBufferedWritableByteChannelTest {
     assertThat(actual).isEqualTo(expected);
   }
 
+  @Property
+  void bufferAllocationShouldOnlyHappenWhenNeeded(@ForAll("BufferSizes") WriteOps writeOps)
+      throws IOException {
+    AuditingBufferHandle handle =
+        new AuditingBufferHandle(BufferHandle.allocate(writeOps.bufferSize));
+    try (ByteArrayOutputStream baos = new ByteArrayOutputStream();
+        CountingWritableByteChannelAdapter adapter =
+            new CountingWritableByteChannelAdapter(Channels.newChannel(baos));
+        BufferedWritableByteChannel c = new DefaultBufferedWritableByteChannel(handle, adapter)) {
+
+      for (ByteBuffer buf : writeOps.writes) {
+        c.write(buf);
+      }
+    }
+
+    // if our write size is evenly divisible by our buffer size AND our buffer size is smaller
+    // than the total data size we expect to never allocate a buffer
+    if (writeOps.writeSize % writeOps.bufferSize == 0
+        && writeOps.bufferSize <= writeOps.bytes.length) {
+      assertThat(handle.getCallCount).isEqualTo(0);
+    } else {
+      assertThat(handle.getCallCount).isGreaterThan(0);
+    }
+  }
+
+  @Provide("BufferSizes")
+  static Arbitrary<WriteOps> arbitraryBufferSizes() {
+    return Arbitraries.of(
+        // expect no allocation
+        WriteOps.of(32, 4, 16),
+        WriteOps.of(32, 16, 16),
+        WriteOps.of(32, 32, 32),
+        // expect allocation
+        WriteOps.of(32, 33, 32),
+        WriteOps.of(32, 64, 4));
+  }
+
   private static final class WriteOps {
     private final byte[] bytes;
     private final int bufferSize;
@@ -397,7 +433,13 @@ public final class DefaultBufferedWritableByteChannelTest {
 
     @Override
     public String toString() {
-      return "WriteOps{"
+      return "[WriteOps.of("
+          + fmt(bytes.length)
+          + ", "
+          + fmt(bufferSize)
+          + ", "
+          + fmt(writeSize)
+          + ")] WriteOps{"
           + "bytes.length="
           + fmt(bytes.length)
           + ", bufferSize="
@@ -562,4 +604,35 @@ public final class DefaultBufferedWritableByteChannelTest {
       }
     }
   }
+
+  private static final class AuditingBufferHandle extends BufferHandle {
+    private final BufferHandle delegate;
+
+    private int getCallCount = 0;
+
+    AuditingBufferHandle(BufferHandle delegate) {
+      this.delegate = delegate;
+    }
+
+    @Override
+    public int remaining() {
+      return delegate.remaining();
+    }
+
+    @Override
+    public int capacity() {
+      return delegate.capacity();
+    }
+
+    @Override
+    public int position() {
+      return delegate.position();
+    }
+
+    @Override
+    public ByteBuffer get() {
+      getCallCount++;
+      return delegate.get();
+    }
+  }
 }