Merge branch 'cassandra-4.1' into cassandra-5.0

* cassandra-4.1:
  Fix memory leak in BufferPoolAllocator when a capacity needs to be extended

src/java/org/apache/cassandra/net/BufferPoolAllocator.java

@@ -91,6 +91,12 @@ public abstract class BufferPoolAllocator extends AbstractByteBufAllocator
         return bufferPool.usedSizeInBytes();
     }
 
+    @VisibleForTesting
+    long overflowMemoryInBytes()
+    {
+        return bufferPool.overflowMemoryInBytes();
+    }
+
     void release()
     {
     }
@@ -117,6 +123,7 @@ public abstract class BufferPoolAllocator extends AbstractByteBufAllocator
 
             ByteBuf newBuffer = super.capacity(newCapacity);
             ByteBuffer nioBuffer = newBuffer.nioBuffer(0, newBuffer.capacity());
+            nioBuffer = bufferPool.unwrapBufferPoolManagedBuffer(nioBuffer);
 
             bufferPool.put(wrapped);
             wrapped = nioBuffer;

src/java/org/apache/cassandra/utils/memory/BufferPool.java

@@ -1648,4 +1648,30 @@ public class BufferPool
                  + (pool.chunks.chunk1 != null ? 1 : 0)
                  + (pool.chunks.chunk2 != null ? 1 : 0);
     }
+
+    /**
+     * @return the inner buffer if it has a BufferPool.Chunk attached
+     *  and originalBuffer in other cases
+     */
+    public ByteBuffer unwrapBufferPoolManagedBuffer(ByteBuffer originalBuffer)
+    {
+        int MAX_DEPTH = 32; // a protection against possible loops in attachments
+        int depth = 0;
+        ByteBuffer buffer = originalBuffer;
+        do
+        {
+            if (buffer == null || !isExactlyDirect(buffer))
+                return originalBuffer;
+            if (Chunk.getParentChunk(buffer) != null)
+                return buffer;
+
+            Object attachment = MemoryUtil.getAttachment(buffer);
+            if (!(attachment instanceof ByteBuffer))
+                return originalBuffer;
+            buffer = (ByteBuffer) attachment;
+            depth++;
+        }
+        while (depth < MAX_DEPTH);
+        return originalBuffer;
+    }
 }

test/unit/org/apache/cassandra/net/BufferPoolAllocatorTest.java

@@ -22,31 +22,44 @@ import java.nio.ByteBuffer;
 import java.util.Arrays;
 import java.util.Random;
 
+import org.junit.BeforeClass;
 import org.junit.Test;
 
 import io.netty.buffer.ByteBuf;
+import org.apache.cassandra.config.DataStorageSpec;
 import org.apache.cassandra.config.DatabaseDescriptor;
+import org.assertj.core.api.Assertions;
 
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
 
 public class BufferPoolAllocatorTest
 {
+
+    @BeforeClass
+    public static void beforeClass()
+    {
+        DatabaseDescriptor.clientInitialization();
+        // cache size hould be more than a macro chunk size for proper pool testing
+        // if it is 0 or less than a macro chunk size we actually do not pool
+        DatabaseDescriptor.getRawConfig().networking_cache_size = new DataStorageSpec.IntMebibytesBound(128);
+    }
+
     @Test
     public void testAdoptedBufferContentAfterResize() {
-        DatabaseDescriptor.clientInitialization();
-        ByteBuf buffer = GlobalBufferPoolAllocator.instance.buffer(200, 500);
-        assertEquals(200, GlobalBufferPoolAllocator.instance.usedSizeInBytes());
-
+        ByteBuf buffer = allocateByteBuf(200, 500);
+        int originalCapacity = buffer.capacity();
         byte[] content = new byte[300];
 
         Random rand = new Random();
         rand.nextBytes(content);
 
         buffer.writeBytes(Arrays.copyOfRange(content, 0, 200));
-        assertEquals(200, GlobalBufferPoolAllocator.instance.usedSizeInBytes());
+        assertEquals(originalCapacity, GlobalBufferPoolAllocator.instance.usedSizeInBytes());
 
         buffer.writeBytes(Arrays.copyOfRange(content, 200, 300));
+        int increasedCapacity = buffer.capacity();
+        assertEquals(increasedCapacity, GlobalBufferPoolAllocator.instance.usedSizeInBytes());
 
         byte[] bufferContent = new byte[300];
 
@@ -54,17 +67,16 @@ public class BufferPoolAllocatorTest
         ByteBuffer adopted = wrapped.adopt();
         adopted.get(bufferContent);
         assertArrayEquals(content, bufferContent);
-        assertEquals(500, GlobalBufferPoolAllocator.instance.usedSizeInBytes());
+        assertEquals(increasedCapacity, GlobalBufferPoolAllocator.instance.usedSizeInBytes());
 
         GlobalBufferPoolAllocator.instance.put(adopted);
-        assertEquals(0, GlobalBufferPoolAllocator.instance.usedSizeInBytes());
+        ensureThatAllMemoryIsReturnedBackToBufferPool();
     }
 
     @Test
     public void testAdoptedBufferContentBeforeResize() {
-        DatabaseDescriptor.clientInitialization();
-        ByteBuf buffer = GlobalBufferPoolAllocator.instance.buffer(200, 300);
-        assertEquals(200, GlobalBufferPoolAllocator.instance.usedSizeInBytes());
+        ByteBuf buffer = allocateByteBuf(200, 300);
+        int originalCapacity = buffer.capacity();
 
         byte[] content = new byte[200];
 
@@ -72,7 +84,7 @@ public class BufferPoolAllocatorTest
         rand.nextBytes(content);
 
         buffer.writeBytes(content);
-        assertEquals(200, GlobalBufferPoolAllocator.instance.usedSizeInBytes());
+        assertEquals(originalCapacity, GlobalBufferPoolAllocator.instance.usedSizeInBytes());
 
         byte[] bufferContent = new byte[200];
 
@@ -82,115 +94,121 @@ public class BufferPoolAllocatorTest
         assertArrayEquals(content, bufferContent);
 
         GlobalBufferPoolAllocator.instance.put(adopted);
-        assertEquals(0, GlobalBufferPoolAllocator.instance.usedSizeInBytes());
+        ensureThatAllMemoryIsReturnedBackToBufferPool();
     }
 
     @Test
     public void testPutPooledBufferBackIntoPool() {
-        DatabaseDescriptor.clientInitialization();
-        ByteBuf buffer = GlobalBufferPoolAllocator.instance.buffer(200, 500);
-        assertEquals(200, GlobalBufferPoolAllocator.instance.usedSizeInBytes());
+        ByteBuf buffer = allocateByteBuf(200, 500);
         buffer.writeBytes(new byte[200]);
 
         buffer.release();
-        assertEquals(0, GlobalBufferPoolAllocator.instance.usedSizeInBytes());
+        ensureThatAllMemoryIsReturnedBackToBufferPool();
     }
 
     @Test
     public void testPutResizedBufferBackIntoPool() {
-        DatabaseDescriptor.clientInitialization();
-        ByteBuf buffer = GlobalBufferPoolAllocator.instance.buffer(200, 500);
-        assertEquals(200, GlobalBufferPoolAllocator.instance.usedSizeInBytes());
+        ByteBuf buffer = allocateByteBuf(200, 500);
         buffer.writeBytes(new byte[500]);
 
         buffer.release();
-        assertEquals(0, GlobalBufferPoolAllocator.instance.usedSizeInBytes());
+        ensureThatAllMemoryIsReturnedBackToBufferPool();
     }
 
     @Test
     public void testBufferDefaultMaxCapacity()
     {
-        DatabaseDescriptor.clientInitialization();
         ByteBuf noMaxCapacity = GlobalBufferPoolAllocator.instance.buffer(100);
         noMaxCapacity.writeBytes(new byte[100]);
         assertEquals(100, noMaxCapacity.readableBytes());
         noMaxCapacity.release();
-        assertEquals(0, GlobalBufferPoolAllocator.instance.usedSizeInBytes());
+        ensureThatAllMemoryIsReturnedBackToBufferPool();
     }
 
     @Test
     public void testBufferWithMaxCapacity()
     {
-        DatabaseDescriptor.clientInitialization();
-        ByteBuf buffer = GlobalBufferPoolAllocator.instance.buffer(100, 500);
+        ByteBuf buffer = allocateByteBuf(100, 500);
         buffer.writeBytes(new byte[500]);
         assertEquals(500, buffer.readableBytes());
-        assertEquals(500, GlobalBufferPoolAllocator.instance.usedSizeInBytes());
+        assertEquals(buffer.capacity(), GlobalBufferPoolAllocator.instance.usedSizeInBytes());
         buffer.release();
-        assertEquals(0, GlobalBufferPoolAllocator.instance.usedSizeInBytes());
+        ensureThatAllMemoryIsReturnedBackToBufferPool();
     }
 
     @Test
     public void testBufferContentAfterResize()
     {
-        DatabaseDescriptor.clientInitialization();
-        ByteBuf buffer = GlobalBufferPoolAllocator.instance.buffer(200, 300);
-        assertEquals(200, GlobalBufferPoolAllocator.instance.usedSizeInBytes());
+        ByteBuf buffer = allocateByteBuf(200, 300);
+        int originalCapacity = buffer.capacity();
 
         byte[] content = new byte[300];
-
         Random rand = new Random();
         rand.nextBytes(content);
 
         buffer.writeBytes(Arrays.copyOfRange(content, 0, 200));
-        assertEquals(200, GlobalBufferPoolAllocator.instance.usedSizeInBytes());
+        assertEquals(originalCapacity, GlobalBufferPoolAllocator.instance.usedSizeInBytes());
 
         buffer.writeBytes(Arrays.copyOfRange(content, 200, 300));
 
         byte[] bufferContent = new byte[300];
         buffer.readBytes(bufferContent);
         assertArrayEquals(content, bufferContent);
-        assertEquals(300, GlobalBufferPoolAllocator.instance.usedSizeInBytes());
+        Assertions.assertThat(buffer.capacity()).isGreaterThanOrEqualTo(300);
+        assertEquals(buffer.capacity(), GlobalBufferPoolAllocator.instance.usedSizeInBytes());
+
         buffer.release();
-        assertEquals(0, GlobalBufferPoolAllocator.instance.usedSizeInBytes());
+        ensureThatAllMemoryIsReturnedBackToBufferPool();
+
     }
 
     @Test(expected = IndexOutOfBoundsException.class)
     public void testBufferExceedMaxCapacity()
     {
-        DatabaseDescriptor.clientInitialization();
-        ByteBuf maxCapacity = GlobalBufferPoolAllocator.instance.buffer(100, 200);
+        ByteBuf maxCapacity = allocateByteBuf(100, 200);
         try
         {
             maxCapacity.writeBytes(new byte[300]);
         } finally {
             maxCapacity.release();
-            assertEquals(0, GlobalBufferPoolAllocator.instance.usedSizeInBytes());
+            ensureThatAllMemoryIsReturnedBackToBufferPool();
         }
     }
 
     @Test
     public void testResizeBufferMultipleTimes()
     {
-        DatabaseDescriptor.clientInitialization();
-        ByteBuf buffer = GlobalBufferPoolAllocator.instance.buffer(100, 2000);
+        ByteBuf buffer = allocateByteBuf(100, 2000);
         buffer.writeBytes(new byte[200]);
         assertEquals(200, buffer.readableBytes());
-        assertEquals(256, buffer.capacity());
-        assertEquals(256, GlobalBufferPoolAllocator.instance.usedSizeInBytes());
+        assertEquals(buffer.capacity(), GlobalBufferPoolAllocator.instance.usedSizeInBytes());
 
         buffer.writeBytes(new byte[100]);
         assertEquals(300, buffer.readableBytes());
-        assertEquals(512, buffer.capacity());
-        assertEquals(512, GlobalBufferPoolAllocator.instance.usedSizeInBytes());
+        assertEquals(buffer.capacity(), GlobalBufferPoolAllocator.instance.usedSizeInBytes());
 
         buffer.writeBytes(new byte[300]);
         assertEquals(600, buffer.readableBytes());
-        assertEquals(1024, buffer.capacity());
-        assertEquals(1024, GlobalBufferPoolAllocator.instance.usedSizeInBytes());
+        assertEquals(buffer.capacity(), GlobalBufferPoolAllocator.instance.usedSizeInBytes());
 
         buffer.release();
-        assertEquals(0, GlobalBufferPoolAllocator.instance.usedSizeInBytes());
+        ensureThatAllMemoryIsReturnedBackToBufferPool();
     }
 
+    private static ByteBuf allocateByteBuf(int initialCapacity, int maxCapacity)
+    {
+        ByteBuf buffer = GlobalBufferPoolAllocator.instance.buffer(initialCapacity, maxCapacity);
+        int originalCapacity = buffer.capacity();
+
+        // BufferPool can allocate more capacity than requested to avoid fragmentation
+        Assertions.assertThat(originalCapacity).isGreaterThanOrEqualTo(initialCapacity);
+        assertEquals(originalCapacity, GlobalBufferPoolAllocator.instance.usedSizeInBytes());
+        return buffer;
+    }
+
+    private static void ensureThatAllMemoryIsReturnedBackToBufferPool()
+    {
+        assertEquals(0, GlobalBufferPoolAllocator.instance.usedSizeInBytes());
+        assertEquals(0, GlobalBufferPoolAllocator.instance.overflowMemoryInBytes());
+    }
 }