From 3f079b23452a44ab443b4519b3b1301c2b4d3563 Mon Sep 17 00:00:00 2001 From: Dmitry Konstantinov Date: Sun, 20 Jul 2025 14:36:13 +0100 Subject: [PATCH] Fix memory leak in BufferPoolAllocator when a capacity needs to be extended Before the fix during a capacity extension BufferPoolAllocator returned to BufferPool a sliced ByteBuffer wrapper object instead of the originally allocated one, so the ByteBuffer was not recycled by BufferPool Adjust BufferPoolAllocatorTest to test the ByteBuf capacity extension with a real BufferPool behavior Patch by Dmitry Konstantinov; reviewed by Michael Semb Wever for CASSANDRA-20753 --- CHANGES.txt | 1 + .../cassandra/net/BufferPoolAllocator.java | 7 ++ .../cassandra/utils/memory/BufferPool.java | 26 +++++ .../net/BufferPoolAllocatorTest.java | 107 ++++++++++-------- 4 files changed, 96 insertions(+), 45 deletions(-) diff --git a/CHANGES.txt b/CHANGES.txt index 8364602a1a..ad5c3f688c 100644 --- a/CHANGES.txt +++ b/CHANGES.txt @@ -1,4 +1,5 @@ 4.0.20 + * Fix memory leak in BufferPoolAllocator when a capacity needs to be extended (CASSANDRA-20753) * Leveled Compaction doesn't validate maxBytesForLevel when the table is altered/created (CASSANDRA-20570) * Updated dtest-api to 0.0.18 and removed JMX-related classes that now live in the dtest-api (CASSANDRA-20884) diff --git a/src/java/org/apache/cassandra/net/BufferPoolAllocator.java b/src/java/org/apache/cassandra/net/BufferPoolAllocator.java index 8c55133e3f..63e734a998 100644 --- a/src/java/org/apache/cassandra/net/BufferPoolAllocator.java +++ b/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; diff --git a/src/java/org/apache/cassandra/utils/memory/BufferPool.java b/src/java/org/apache/cassandra/utils/memory/BufferPool.java index a92f63782a..73ebd8e868 100644 --- a/src/java/org/apache/cassandra/utils/memory/BufferPool.java +++ b/src/java/org/apache/cassandra/utils/memory/BufferPool.java @@ -1613,4 +1613,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; + } } diff --git a/test/unit/org/apache/cassandra/net/BufferPoolAllocatorTest.java b/test/unit/org/apache/cassandra/net/BufferPoolAllocatorTest.java index 3d902a01d9..a5c4676e1b 100644 --- a/test/unit/org/apache/cassandra/net/BufferPoolAllocatorTest.java +++ b/test/unit/org/apache/cassandra/net/BufferPoolAllocatorTest.java @@ -22,31 +22,43 @@ 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.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_in_mb = 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 +66,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 +83,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 +93,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()); + } }