Reference-counts each chunk: when all sub-allocations are released, the chunk is recycled
@@ -106,7 +109,7 @@
* }
*
* For allocations >= chunkSize, a dedicated memory region is allocated directly (no bump
- * pointer), behaving identically to {@code UnsafeAllocationManager}.
+ * pointer), using a dedicated Netty {@code ByteBuf}.
*/
public class ChunkedAllocationManager extends AllocationManager {
@@ -114,7 +117,7 @@ public class ChunkedAllocationManager extends AllocationManager {
private static final long ALIGNMENT = 8;
/** Default chunk size: 4MB (matches Netty 4.1+ maxOrder=9). */
- private static final long DEFAULT_CHUNK_SIZE = 4L * 1024 * 1024;
+ private static final int DEFAULT_CHUNK_SIZE = 4 * 1024 * 1024;
/** Default maximum number of empty chunks to keep in the free-list. */
private static final int DEFAULT_MAX_FREE_CHUNKS = 3;
@@ -126,7 +129,8 @@ public class ChunkedAllocationManager extends AllocationManager {
private final long offsetInChunk;
// --- Fields for direct allocation (large request, owns its own memory) ---
- private final long directAddress;
+ private final ByteBuf directByteBuf;
+ private final ChunkedFactory chunkedFactory;
/** Sub-allocation carved from a shared {@link Chunk}. */
private ChunkedAllocationManager(
@@ -135,16 +139,22 @@ private ChunkedAllocationManager(
this.chunk = chunk;
this.offsetInChunk = offset;
this.allocatedSize = size;
- this.directAddress = 0;
+ this.directByteBuf = null;
+ this.chunkedFactory = null;
}
/** Direct allocation for oversized requests (>= chunkSize). Owns its own memory region. */
- private ChunkedAllocationManager(BufferAllocator accountingAllocator, long address, long size) {
+ private ChunkedAllocationManager(
+ BufferAllocator accountingAllocator,
+ ByteBuf directByteBuf,
+ long size,
+ ChunkedFactory chunkedFactory) {
super(accountingAllocator);
this.chunk = null;
this.offsetInChunk = 0;
this.allocatedSize = size;
- this.directAddress = address;
+ this.directByteBuf = directByteBuf;
+ this.chunkedFactory = chunkedFactory;
}
@Override
@@ -155,9 +165,9 @@ public long getSize() {
@Override
protected long memoryAddress() {
if (chunk != null) {
- return chunk.address + offsetInChunk;
+ return chunk.directByteBuf.memoryAddress() + offsetInChunk;
}
- return directAddress;
+ return directByteBuf.memoryAddress();
}
@Override
@@ -165,7 +175,8 @@ protected void release0() {
if (chunk != null) {
chunk.releaseSubAllocation();
} else {
- MemoryUtil.UNSAFE.freeMemory(directAddress);
+ directByteBuf.release();
+ chunkedFactory.decrementDirectMemoryBytes(allocatedSize);
}
}
@@ -174,12 +185,12 @@ protected void release0() {
// -------------------------------------------------------------------------
/**
- * A contiguous native memory region that holds multiple small allocations via bump-pointer.
+ * A contiguous direct memory region that holds multiple small allocations via bump-pointer.
* Reference-counted: when all sub-allocations are released (count reaches 0), the chunk is
* recycled back to the factory's free-list.
*/
static class Chunk {
- final long address;
+ final ByteBuf directByteBuf;
final long capacity;
/** Bump pointer — only accessed under the factory's synchronized lock. */
long used;
@@ -198,8 +209,8 @@ static class Chunk {
/** Back-reference to the owning factory for recycling on drain. */
final ChunkedFactory factory;
- Chunk(long capacity, ChunkedFactory factory) {
- this.address = MemoryUtil.UNSAFE.allocateMemory(capacity);
+ Chunk(int capacity, ChunkedFactory factory) {
+ this.directByteBuf = UnpooledByteBufAllocator.DEFAULT.directBuffer(capacity);
this.capacity = capacity;
this.used = 0;
this.factory = factory;
@@ -266,9 +277,9 @@ void resetBump() {
// subAllocCount is already 0 at this point.
}
- /** Frees the underlying native memory. */
+ /** Deterministically releases the underlying direct memory. */
void destroy() {
- MemoryUtil.UNSAFE.freeMemory(address);
+ directByteBuf.release();
}
}
@@ -288,7 +299,7 @@ void destroy() {
*/
public static class ChunkedFactory implements AllocationManager.Factory {
- private final long chunkSize;
+ private final int chunkSize;
private final int maxFreeChunks;
/** The chunk currently receiving bump allocations. May be null initially. */
@@ -297,6 +308,9 @@ public static class ChunkedFactory implements AllocationManager.Factory {
/** Pool of empty chunks available for reuse. */
private final Deque freeChunks = new ArrayDeque<>();
+ /** Total direct memory, in bytes, currently allocated by this factory. */
+ private final AtomicLong directMemoryAllocatedBytes = new AtomicLong();
+
/** Set to true when {@link #close()} is called. */
private boolean closed;
@@ -311,7 +325,7 @@ public ChunkedFactory() {
* @param chunkSize maximum size of each chunk (bytes). Allocations >= this go direct.
* @param maxFreeChunks maximum number of empty chunks to keep cached for reuse.
*/
- public ChunkedFactory(long chunkSize, int maxFreeChunks) {
+ public ChunkedFactory(int chunkSize, int maxFreeChunks) {
this.chunkSize = chunkSize;
this.maxFreeChunks = maxFreeChunks;
}
@@ -319,10 +333,18 @@ public ChunkedFactory(long chunkSize, int maxFreeChunks) {
@Override
public synchronized AllocationManager create(
BufferAllocator accountingAllocator, long size) {
+ if (closed) {
+ throw new IllegalStateException("ChunkedFactory has been closed.");
+ }
if (size > chunkSize) {
- // Large allocation: give it its own memory region.
- long address = MemoryUtil.UNSAFE.allocateMemory(size);
- return new ChunkedAllocationManager(accountingAllocator, address, size);
+ // Large allocation: use Netty direct buffer for deterministic release.
+ if (size > Integer.MAX_VALUE) {
+ throw new IllegalArgumentException(
+ "Allocation size " + size + " exceeds maximum " + Integer.MAX_VALUE);
+ }
+ ByteBuf directByteBuf = UnpooledByteBufAllocator.DEFAULT.directBuffer((int) size);
+ directMemoryAllocatedBytes.addAndGet(size);
+ return new ChunkedAllocationManager(accountingAllocator, directByteBuf, size, this);
}
// Align to 8 bytes for safe direct-memory access.
@@ -352,7 +374,9 @@ private synchronized Chunk obtainChunk() {
recycled.resetBump();
return recycled;
}
- return new Chunk(chunkSize, this);
+ Chunk chunk = new Chunk(chunkSize, this);
+ directMemoryAllocatedBytes.addAndGet(chunkSize);
+ return chunk;
}
/**
@@ -381,7 +405,7 @@ synchronized void onChunkDrained(Chunk chunk, int expectedGeneration) {
if (closed) {
// Factory is closed — no one will ever reuse this chunk. Free it.
- chunk.destroy();
+ destroyChunk(chunk);
} else if (chunk == activeChunk) {
// Still the active chunk — just reset bump pointer for continued use.
chunk.resetBump();
@@ -389,8 +413,8 @@ synchronized void onChunkDrained(Chunk chunk, int expectedGeneration) {
// Not active, pool has room — recycle.
freeChunks.offerFirst(chunk);
} else {
- // Pool is full — free native memory.
- chunk.destroy();
+ // Pool is full — free direct memory.
+ destroyChunk(chunk);
}
}
@@ -403,12 +427,31 @@ public synchronized void close() {
closed = true;
while (!freeChunks.isEmpty()) {
Chunk poll = freeChunks.poll();
- poll.destroy();
- }
- if (activeChunk != null && activeChunk.subAllocCount.get() == 0) {
- activeChunk.destroy();
+ destroyChunk(poll);
}
+ Chunk chunk = activeChunk;
activeChunk = null;
+ if (chunk != null && chunk.subAllocCount.get() == 0) {
+ // Invalidate a pending onChunkDrained callback that observed this zero count before
+ // close acquired the factory lock.
+ chunk.drainGeneration++;
+ destroyChunk(chunk);
+ }
+ }
+
+ /** Returns the direct memory, in bytes, currently allocated by this factory. */
+ public long getDirectMemoryAllocatedBytes() {
+ return directMemoryAllocatedBytes.get();
+ }
+
+ private void destroyChunk(Chunk chunk) {
+ long cap = chunk.capacity;
+ chunk.destroy();
+ directMemoryAllocatedBytes.addAndGet(-cap);
+ }
+
+ void decrementDirectMemoryBytes(long size) {
+ directMemoryAllocatedBytes.addAndGet(-size);
}
@VisibleForTesting
diff --git a/fluss-common/src/test/java/org/apache/fluss/shaded/arrow/org/apache/arrow/memory/ChunkedAllocationManagerTest.java b/fluss-common/src/test/java/org/apache/fluss/shaded/arrow/org/apache/arrow/memory/ChunkedAllocationManagerTest.java
index e6c1bc01f90..a9417e2f8d9 100644
--- a/fluss-common/src/test/java/org/apache/fluss/shaded/arrow/org/apache/arrow/memory/ChunkedAllocationManagerTest.java
+++ b/fluss-common/src/test/java/org/apache/fluss/shaded/arrow/org/apache/arrow/memory/ChunkedAllocationManagerTest.java
@@ -29,12 +29,13 @@
import java.util.concurrent.atomic.AtomicReference;
import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
/** Tests for {@link ChunkedAllocationManager}. */
class ChunkedAllocationManagerTest {
/** Use a small chunk size (1KB) for testing to make chunk transitions easy to trigger. */
- private static final long TEST_CHUNK_SIZE = 1024;
+ private static final int TEST_CHUNK_SIZE = 1024;
private static final int TEST_MAX_FREE_CHUNKS = 2;
@@ -70,6 +71,45 @@ void testSmallAllocationPackedIntoSameChunk() {
buf3.close();
}
+ @Test
+ void testDirectMemoryAllocatedBytesTracksChunksAndDirectAllocations() {
+ assertThat(factory.getDirectMemoryAllocatedBytes()).isZero();
+
+ try (ArrowBuf small = allocator.buffer(64)) {
+ assertThat(small.capacity()).isEqualTo(64L);
+ assertThat(factory.getDirectMemoryAllocatedBytes()).isEqualTo(TEST_CHUNK_SIZE);
+
+ try (ArrowBuf direct = allocator.buffer(TEST_CHUNK_SIZE + 1)) {
+ assertThat(direct.capacity()).isEqualTo(2 * TEST_CHUNK_SIZE);
+ // allocate a new direct buffer(which not resued)
+ assertThat(factory.getDirectMemoryAllocatedBytes())
+ .isEqualTo(TEST_CHUNK_SIZE + direct.capacity());
+ }
+
+ assertThat(factory.getDirectMemoryAllocatedBytes()).isEqualTo(TEST_CHUNK_SIZE);
+
+ try (ArrowBuf anotherSmall = allocator.buffer(64)) {
+ assertThat(anotherSmall.capacity()).isEqualTo(64L);
+ // reuse the same chunk.
+ assertThat(factory.getDirectMemoryAllocatedBytes()).isEqualTo(TEST_CHUNK_SIZE);
+ }
+ }
+
+ // The empty active chunk remains cached until the factory is closed.
+ assertThat(factory.getDirectMemoryAllocatedBytes()).isEqualTo(TEST_CHUNK_SIZE);
+ factory.close();
+ assertThat(factory.getDirectMemoryAllocatedBytes()).isZero();
+ }
+
+ @Test
+ void testAllocationRejectedAfterFactoryClose() {
+ factory.close();
+
+ assertThatThrownBy(() -> allocator.buffer(64))
+ .isInstanceOf(IllegalStateException.class)
+ .hasMessage("ChunkedFactory has been closed.");
+ }
+
@Test
void testAlignmentBytes() {
// Allocate a non-8-byte-aligned size; the next allocation should still be 8-byte aligned.
@@ -280,6 +320,7 @@ void testCloseWhileArrowBufsAliveDoesNotLeakMemory() {
for (int i = 0; i < 4; i++) {
bufs.add(localAllocator.buffer(64));
}
+ assertThat(localFactory.getDirectMemoryAllocatedBytes()).isEqualTo(TEST_CHUNK_SIZE);
// Close factory while ArrowBufs are still alive.
@@ -287,12 +328,74 @@ void testCloseWhileArrowBufsAliveDoesNotLeakMemory() {
// Active chunk was not destroyed (subAllocCount > 0), just nulled.
assertThat(localFactory.getActiveChunk()).isNull();
+ assertThat(localFactory.getDirectMemoryAllocatedBytes()).isEqualTo(TEST_CHUNK_SIZE);
// Release all ArrowBufs — onChunkDrained should destroy the chunk, not recycle it.
bufs.forEach(ArrowBuf::close);
// The chunk must NOT have been added to freeChunks (would be a leak).
assertThat(localFactory.getFreeChunks()).isEmpty();
+ assertThat(localFactory.getDirectMemoryAllocatedBytes()).isZero();
+ localAllocator.close();
+ }
+
+ @Test
+ void testDirectAllocationReleaseAfterFactoryClose() {
+ // Verify that closing the factory while a direct (large) ArrowBuf is alive
+ // does not leak memory, and the counter correctly drops to zero after release.
+ ChunkedFactory localFactory = new ChunkedFactory(TEST_CHUNK_SIZE, TEST_MAX_FREE_CHUNKS);
+ BufferAllocator localAllocator = BufferAllocatorUtil.createBufferAllocator(localFactory);
+
+ // Allocate a large buffer that bypasses chunk sub-allocation.
+ ArrowBuf directBuf = localAllocator.buffer(TEST_CHUNK_SIZE + 1);
+ long directCapacity = directBuf.capacity();
+ assertThat(localFactory.getDirectMemoryAllocatedBytes()).isEqualTo(directCapacity);
+
+ // Close factory while the direct ArrowBuf is still alive.
+ localFactory.close();
+
+ // Memory is still tracked (ArrowBuf not yet released).
+ assertThat(localFactory.getDirectMemoryAllocatedBytes()).isEqualTo(directCapacity);
+
+ // Release the ArrowBuf — should properly free the direct ByteBuf and decrement.
+ directBuf.close();
+
+ // The direct allocation is freed.
+ assertThat(localFactory.getDirectMemoryAllocatedBytes()).isZero();
+ localAllocator.close();
+ }
+
+ @Test
+ void testDirectAllocationExceedsMaxIntSize() {
+ ChunkedFactory localFactory = new ChunkedFactory(TEST_CHUNK_SIZE, TEST_MAX_FREE_CHUNKS);
+ BufferAllocator localAllocator = BufferAllocatorUtil.createBufferAllocator(localFactory);
+
+ // Request a size larger than Integer.MAX_VALUE — should be rejected.
+ assertThatThrownBy(() -> localAllocator.buffer(Integer.MAX_VALUE + 1L))
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessageContaining("Allocation size");
+
+ localAllocator.close();
+ localFactory.close();
+ }
+
+ @Test
+ void testDirectAllocationMultipleCycles() {
+ // Verify that multiple direct allocate-release cycles keep the counter accurate.
+ ChunkedFactory localFactory = new ChunkedFactory(TEST_CHUNK_SIZE, TEST_MAX_FREE_CHUNKS);
+ BufferAllocator localAllocator = BufferAllocatorUtil.createBufferAllocator(localFactory);
+
+ long directSize = TEST_CHUNK_SIZE + 1;
+ for (int cycle = 0; cycle < 5; cycle++) {
+ ArrowBuf directBuf = localAllocator.buffer(directSize);
+ long directCapacity = directBuf.capacity();
+ assertThat(localFactory.getDirectMemoryAllocatedBytes()).isEqualTo(directCapacity);
+ directBuf.close();
+ assertThat(localFactory.getDirectMemoryAllocatedBytes()).isZero();
+ }
+
+ localFactory.close();
+ assertThat(localFactory.getDirectMemoryAllocatedBytes()).isZero();
localAllocator.close();
}
}