Introduce ByteBuf.maxFastWritableBytes() method (#9086)

Motivation

ByteBuf capacity is automatically increased as needed up to maxCapacity
when writing beyond the buffer's current capacity. However there's no
way to tell in general whether such an increase will result in a
relatively costly internal buffer re-allocation.

For unpooled buffers it always does, in pooled cases it depends on the
size of the associated chunk of allocated memory, which I don't think is
currently exposed in any way.

It would sometimes be useful to know where this limit is when making
external decisions about whether to reuse or preemptively reallocate.

It would also be advantageous to take this limit into account when
auto-increasing the capacity during writes, to defer such reallocation
until really necessary.

Modifications

Introduce new AbstractByteBuf.maxFastWritableBytes() method which will
return a value >= writableBytes() and <= maxWritableBytes().

Make use of the new method in the sizing decision made by the
AbstractByteBuf.ensureWritable(...) methods.

Result

Less reallocation/copying.

buffer/src/main/java/io/netty/buffer/AbstractByteBuf.java

@@ -230,6 +230,7 @@ public abstract class AbstractByteBuf extends ByteBuf {
         if (minWritableBytes <= writableBytes()) {
             return;
         }
+        final int writerIndex = writerIndex();
         if (checkBounds) {
             if (minWritableBytes > maxCapacity - writerIndex) {
                 throw new IndexOutOfBoundsException(String.format(
@@ -239,7 +240,14 @@ public abstract class AbstractByteBuf extends ByteBuf {
         }
 
         // Normalize the current capacity to the power of 2.
-        int newCapacity = alloc().calculateNewCapacity(writerIndex + minWritableBytes, maxCapacity);
+        int minNewCapacity = writerIndex + minWritableBytes;
+        int newCapacity = alloc().calculateNewCapacity(minNewCapacity, maxCapacity);
+
+        int fastCapacity = writerIndex + maxFastWritableBytes();
+        // Grow by a smaller amount if it will avoid reallocation
+        if (newCapacity > fastCapacity && minNewCapacity <= fastCapacity) {
+            newCapacity = fastCapacity;
+        }
 
         // Adjust to the new capacity.
         capacity(newCapacity);
@@ -266,7 +274,14 @@ public abstract class AbstractByteBuf extends ByteBuf {
         }
 
         // Normalize the current capacity to the power of 2.
-        int newCapacity = alloc().calculateNewCapacity(writerIndex + minWritableBytes, maxCapacity);
+        int minNewCapacity = writerIndex + minWritableBytes;
+        int newCapacity = alloc().calculateNewCapacity(minNewCapacity, maxCapacity);
+
+        int fastCapacity = writerIndex + maxFastWritableBytes();
+        // Grow by a smaller amount if it will avoid reallocation
+        if (newCapacity > fastCapacity && minNewCapacity <= fastCapacity) {
+            newCapacity = fastCapacity;
+        }
 
         // Adjust to the new capacity.
         capacity(newCapacity);

buffer/src/main/java/io/netty/buffer/ByteBuf.java

@@ -236,7 +236,6 @@ import java.nio.charset.UnsupportedCharsetException;
  * Please refer to {@link ByteBufInputStream} and
  * {@link ByteBufOutputStream}.
  */
-@SuppressWarnings("ClassMayBeInterface")
 public abstract class ByteBuf implements ReferenceCounted, Comparable<ByteBuf> {
 
     /**
@@ -413,6 +412,15 @@ public abstract class ByteBuf implements ReferenceCounted, Comparable<ByteBuf> {
      */
     public abstract int maxWritableBytes();
 
+    /**
+     * Returns the maximum number of bytes which can be written for certain without involving
+     * an internal reallocation or data-copy. The returned value will be &ge; {@link #writableBytes()}
+     * and &le; {@link #maxWritableBytes()}.
+     */
+    public int maxFastWritableBytes() {
+        return writableBytes();
+    }
+
     /**
      * Returns {@code true}
      * if and only if {@code (this.writerIndex - this.readerIndex)} is greater

buffer/src/main/java/io/netty/buffer/PooledByteBuf.java

@@ -81,6 +81,11 @@ abstract class PooledByteBuf<T> extends AbstractReferenceCountedByteBuf {
         return length;
     }
 
+    @Override
+    public int maxFastWritableBytes() {
+        return Math.min(maxLength, maxCapacity()) - writerIndex;
+    }
+
     @Override
     public final ByteBuf capacity(int newCapacity) {
         checkNewCapacity(newCapacity);

buffer/src/main/java/io/netty/buffer/SwappedByteBuf.java

@@ -149,6 +149,11 @@ public class SwappedByteBuf extends ByteBuf {
         return buf.maxWritableBytes();
     }
 
+    @Override
+    public int maxFastWritableBytes() {
+        return buf.maxFastWritableBytes();
+    }
+
     @Override
     public boolean isReadable() {
         return buf.isReadable();

buffer/src/main/java/io/netty/buffer/WrappedByteBuf.java

@@ -151,6 +151,11 @@ class WrappedByteBuf extends ByteBuf {
         return buf.maxWritableBytes();
     }
 
+    @Override
+    public int maxFastWritableBytes() {
+        return buf.maxFastWritableBytes();
+    }
+
     @Override
     public final boolean isReadable() {
         return buf.isReadable();

buffer/src/main/java/io/netty/buffer/WrappedCompositeByteBuf.java

@@ -98,6 +98,11 @@ class WrappedCompositeByteBuf extends CompositeByteBuf {
         return wrapped.maxWritableBytes();
     }
 
+    @Override
+    public int maxFastWritableBytes() {
+        return wrapped.maxFastWritableBytes();
+    }
+
     @Override
     public int ensureWritable(int minWritableBytes, boolean force) {
         return wrapped.ensureWritable(minWritableBytes, force);

buffer/src/test/java/io/netty/buffer/AbstractByteBufTest.java

@@ -4846,4 +4846,16 @@ public abstract class AbstractByteBufTest {
             buffer.release();
         }
     }
+
+    @Test
+    public void testMaxFastWritableBytes() {
+        ByteBuf buffer = newBuffer(150, 500).writerIndex(100);
+        assertEquals(50, buffer.writableBytes());
+        assertEquals(150, buffer.capacity());
+        assertEquals(500, buffer.maxCapacity());
+        assertEquals(400, buffer.maxWritableBytes());
+        // Default implementation has fast writable == writable
+        assertEquals(50, buffer.maxFastWritableBytes());
+        buffer.release();
+    }
 }

buffer/src/test/java/io/netty/buffer/AbstractPooledByteBufTest.java

@@ -21,6 +21,8 @@ import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.greaterThanOrEqualTo;
 import static org.hamcrest.Matchers.is;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotEquals;
+import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
 public abstract class AbstractPooledByteBufTest extends AbstractByteBufTest {
@@ -55,4 +57,61 @@ public abstract class AbstractPooledByteBufTest extends AbstractByteBufTest {
             buf.release();
         }
     }
+
+    @Override
+    @Test
+    public void testMaxFastWritableBytes() {
+        ByteBuf buffer = newBuffer(150, 500).writerIndex(100);
+        assertEquals(50, buffer.writableBytes());
+        assertEquals(150, buffer.capacity());
+        assertEquals(500, buffer.maxCapacity());
+        assertEquals(400, buffer.maxWritableBytes());
+
+        int chunkSize = pooledByteBuf(buffer).maxLength;
+        assertTrue(chunkSize >= 150);
+        int remainingInAlloc = Math.min(chunkSize - 100, 400);
+        assertEquals(remainingInAlloc, buffer.maxFastWritableBytes());
+
+        // write up to max, chunk alloc should not change (same handle)
+        long handleBefore = pooledByteBuf(buffer).handle;
+        buffer.writeBytes(new byte[remainingInAlloc]);
+        assertEquals(handleBefore, pooledByteBuf(buffer).handle);
+
+        assertEquals(0, buffer.maxFastWritableBytes());
+        // writing one more should trigger a reallocation (new handle)
+        buffer.writeByte(7);
+        assertNotEquals(handleBefore, pooledByteBuf(buffer).handle);
+
+        // should not exceed maxCapacity even if chunk alloc does
+        buffer.capacity(500);
+        assertEquals(500 - buffer.writerIndex(), buffer.maxFastWritableBytes());
+        buffer.release();
+    }
+
+    private static PooledByteBuf<?> pooledByteBuf(ByteBuf buffer) {
+        // might need to unwrap if swapped (LE) and/or leak-aware-wrapped
+        while (!(buffer instanceof PooledByteBuf)) {
+            buffer = buffer.unwrap();
+        }
+        return (PooledByteBuf<?>) buffer;
+    }
+
+    @Test
+    public void testEnsureWritableDoesntGrowTooMuch() {
+        ByteBuf buffer = newBuffer(150, 500).writerIndex(100);
+
+        assertEquals(50, buffer.writableBytes());
+        int fastWritable = buffer.maxFastWritableBytes();
+        assertTrue(fastWritable > 50);
+
+        long handleBefore = pooledByteBuf(buffer).handle;
+
+        // capacity expansion should not cause reallocation
+        // (should grow precisely the specified amount)
+        buffer.ensureWritable(fastWritable);
+        assertEquals(handleBefore, pooledByteBuf(buffer).handle);
+        assertEquals(100 + fastWritable, buffer.capacity());
+        assertEquals(buffer.writableBytes(), buffer.maxFastWritableBytes());
+        buffer.release();
+    }
 }