From af52a13e891a32cecafa01122187ca1089893548 Mon Sep 17 00:00:00 2001
From: Hongshun Wang <loserwang1024@gmail>
Date: Sun, 20 Sep 2026 17:06:56 +0800
Subject: [PATCH 1/2] [client] Replace Unsafe with Netty direct memory in
 ChunkedAllocationManager and expose metrics.

---
 .../fluss/client/write/RecordAccumulator.java |  11 +-
 .../client/write/RecordAccumulatorTest.java   |  72 +++++++++++-
 .../org/apache/fluss/metrics/MetricNames.java |   8 ++
 .../memory/ChunkedAllocationManager.java      | 103 ++++++++++++-----
 .../memory/ChunkedAllocationManagerTest.java  | 105 +++++++++++++++++-
 5 files changed, 266 insertions(+), 33 deletions(-)

diff --git a/fluss-client/src/main/java/org/apache/fluss/client/write/RecordAccumulator.java b/fluss-client/src/main/java/org/apache/fluss/client/write/RecordAccumulator.java
index 32f4b881e4e..5ff98c64406 100644
--- a/fluss-client/src/main/java/org/apache/fluss/client/write/RecordAccumulator.java
+++ b/fluss-client/src/main/java/org/apache/fluss/client/write/RecordAccumulator.java
@@ -106,7 +106,7 @@ public final class RecordAccumulator {
     /** The arrow buffer allocator to allocate memory for arrow log write batch. */
     private final BufferAllocator bufferAllocator;
 
-    /** The chunked allocation manager factory, stored for explicit native memory release. */
+    /** The chunked allocation manager factory, stored for explicit direct memory release. */
     private final ChunkedAllocationManager.ChunkedFactory chunkedFactory;
 
     /** Coordinates batch memory deallocation with resource destruction. */
@@ -206,6 +206,15 @@ private void registerMetrics(WriterMetricGroup writerMetricGroup) {
         // The number of user threads blocked waiting for buffer memory to enqueue their records
         writerMetricGroup.gauge(
                 MetricNames.WRITER_BUFFER_WAITING_THREADS, writerBufferPool::queued);
+        writerMetricGroup.gauge(
+                MetricNames.WRITER_ACCUMULATOR_HEAP_MEMORY_USED_BYTES,
+                () -> writerBufferPool.totalSize() - writerBufferPool.availableMemory());
+        writerMetricGroup.gauge(
+                MetricNames.WRITER_ACCUMULATOR_ARROW_MEMORY_USED_BYTES,
+                bufferAllocator::getAllocatedMemory);
+        writerMetricGroup.gauge(
+                MetricNames.WRITER_ACCUMULATOR_DIRECT_MEMORY_USED_BYTES,
+                chunkedFactory::getDirectMemoryUsedBytes);
     }
 
     /** Assigns and appends a record using the layout owned by its write context. */
diff --git a/fluss-client/src/test/java/org/apache/fluss/client/write/RecordAccumulatorTest.java b/fluss-client/src/test/java/org/apache/fluss/client/write/RecordAccumulatorTest.java
index 4ea64f4efea..f2d2828bd53 100644
--- a/fluss-client/src/test/java/org/apache/fluss/client/write/RecordAccumulatorTest.java
+++ b/fluss-client/src/test/java/org/apache/fluss/client/write/RecordAccumulatorTest.java
@@ -32,6 +32,8 @@
 import org.apache.fluss.metadata.TableDescriptor;
 import org.apache.fluss.metadata.TableInfo;
 import org.apache.fluss.metadata.TablePath;
+import org.apache.fluss.metrics.Gauge;
+import org.apache.fluss.metrics.MetricNames;
 import org.apache.fluss.record.ChangeType;
 import org.apache.fluss.record.DefaultKvRecord;
 import org.apache.fluss.record.IndexedLogRecord;
@@ -148,6 +150,58 @@ public void start() {
     // TODO Add more tests to test lingMs, retryBackoffMs, deliveryTimeoutMs and
     //  nextBatchExpiryTimeMs if we introduced.
 
+    @Test
+    void testAccumulatorMemoryMetrics() throws Exception {
+        int batchSize = 1024;
+        TestingWriterMetricGroup metrics = TestingWriterMetricGroup.newInstance();
+        RecordAccumulator accum =
+                createTestRecordAccumulator(
+                        5000, batchSize, 256, 10L * batchSize, bucketAssigner, metrics);
+
+        try {
+            assertThat(metrics.getMetrics())
+                    .containsKeys(
+                            MetricNames.WRITER_ACCUMULATOR_HEAP_MEMORY_USED_BYTES,
+                            MetricNames.WRITER_ACCUMULATOR_ARROW_MEMORY_USED_BYTES,
+                            MetricNames.WRITER_ACCUMULATOR_DIRECT_MEMORY_USED_BYTES);
+
+            Gauge<?> heapMemoryUsed =
+                    (Gauge<?>)
+                            metrics.getMetrics()
+                                    .get(MetricNames.WRITER_ACCUMULATOR_HEAP_MEMORY_USED_BYTES);
+            Gauge<?> arrowMemoryUsed =
+                    (Gauge<?>)
+                            metrics.getMetrics()
+                                    .get(MetricNames.WRITER_ACCUMULATOR_ARROW_MEMORY_USED_BYTES);
+            Gauge<?> directMemoryUsedBytes =
+                    (Gauge<?>)
+                            metrics.getMetrics()
+                                    .get(MetricNames.WRITER_ACCUMULATOR_DIRECT_MEMORY_USED_BYTES);
+
+            assertThat(((Number) heapMemoryUsed.getValue()).longValue()).isZero();
+            assertThat(((Number) arrowMemoryUsed.getValue()).longValue()).isZero();
+            assertThat(((Number) directMemoryUsedBytes.getValue()).longValue()).isZero();
+
+            bucketAssigner.setBucketId(0);
+            accum.append(
+                    WriteRecord.forArrowAppend(
+                            DATA1_TABLE_INFO, DATA1_PHYSICAL_TABLE_PATH, row(1, "a"), null),
+                    (bucket, offset, exception) -> {},
+                    cluster);
+
+            long heapMemory = ((Number) heapMemoryUsed.getValue()).longValue();
+            long arrowMemory = ((Number) arrowMemoryUsed.getValue()).longValue();
+            long directMemory = ((Number) directMemoryUsedBytes.getValue()).longValue();
+            assertThat(heapMemory).isPositive();
+            assertThat(arrowMemory).isPositive();
+            assertThat(directMemory).isGreaterThanOrEqualTo(arrowMemory);
+        } finally {
+            accum.close();
+            accum.abortAllBatches(new RuntimeException("test cleanup"));
+            accum.destroyResources();
+        }
+    }
+
     @Test
     void testDrainBatches() throws Exception {
         // test case: node1(tb1, tb2), node2(tb3).
@@ -815,6 +869,22 @@ private RecordAccumulator createTestRecordAccumulator(
             int pageSize,
             long totalSize,
             BucketAssigner assigner) {
+        return createTestRecordAccumulator(
+                batchTimeoutMs,
+                batchSize,
+                pageSize,
+                totalSize,
+                assigner,
+                TestingWriterMetricGroup.newInstance());
+    }
+
+    private RecordAccumulator createTestRecordAccumulator(
+            int batchTimeoutMs,
+            int batchSize,
+            int pageSize,
+            long totalSize,
+            BucketAssigner assigner,
+            TestingWriterMetricGroup metrics) {
         conf.set(ConfigOptions.CLIENT_WRITER_BATCH_TIMEOUT, Duration.ofMillis(batchTimeoutMs));
         // TODO client writer buffer maybe removed.
         conf.set(ConfigOptions.CLIENT_WRITER_BUFFER_MEMORY_SIZE, new MemorySize(totalSize));
@@ -832,7 +902,7 @@ private RecordAccumulator createTestRecordAccumulator(
                                 RpcClient.create(conf, TestingClientMetricGroup.newInstance()),
                                 TabletServerGateway.class),
                         null),
-                TestingWriterMetricGroup.newInstance(),
+                metrics,
                 clock,
                 (tableInfo, path) -> assigner);
     }
diff --git a/fluss-common/src/main/java/org/apache/fluss/metrics/MetricNames.java b/fluss-common/src/main/java/org/apache/fluss/metrics/MetricNames.java
index ac60ee5f101..dd9c844efff 100644
--- a/fluss-common/src/main/java/org/apache/fluss/metrics/MetricNames.java
+++ b/fluss-common/src/main/java/org/apache/fluss/metrics/MetricNames.java
@@ -337,6 +337,14 @@ public class MetricNames {
     public static final String WRITER_RECORDS_PER_BATCH = "recordsPerBatch";
     public static final String WRITER_SEND_LATENCY_MS = "sendLatencyMs";
 
+    // for record accumulator memory
+    public static final String WRITER_ACCUMULATOR_HEAP_MEMORY_USED_BYTES =
+            "accumulatorHeapMemoryUsedBytes";
+    public static final String WRITER_ACCUMULATOR_ARROW_MEMORY_USED_BYTES =
+            "accumulatorArrowMemoryUsedBytes";
+    public static final String WRITER_ACCUMULATOR_DIRECT_MEMORY_USED_BYTES =
+            "accumulatorDirectMemoryUsedBytes";
+
     // for scanner
     public static final String SCANNER_TIME_MS_BETWEEN_POLL = "timeMsBetweenPoll";
     public static final String SCANNER_LAST_POLL_SECONDS_AGO = "lastPollSecondsAgo";
diff --git a/fluss-common/src/main/java/org/apache/fluss/shaded/arrow/org/apache/arrow/memory/ChunkedAllocationManager.java b/fluss-common/src/main/java/org/apache/fluss/shaded/arrow/org/apache/arrow/memory/ChunkedAllocationManager.java
index 14e29a2c696..b032c657bdb 100644
--- a/fluss-common/src/main/java/org/apache/fluss/shaded/arrow/org/apache/arrow/memory/ChunkedAllocationManager.java
+++ b/fluss-common/src/main/java/org/apache/fluss/shaded/arrow/org/apache/arrow/memory/ChunkedAllocationManager.java
@@ -18,11 +18,13 @@
 package org.apache.fluss.shaded.arrow.org.apache.arrow.memory;
 
 import org.apache.fluss.annotation.VisibleForTesting;
-import org.apache.fluss.shaded.arrow.org.apache.arrow.memory.util.MemoryUtil;
+import org.apache.fluss.shaded.netty4.io.netty.buffer.ByteBuf;
+import org.apache.fluss.shaded.netty4.io.netty.buffer.UnpooledByteBufAllocator;
 
 import java.util.ArrayDeque;
 import java.util.Deque;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicLong;
 
 /**
  * An {@link AllocationManager} that packs small allocations into large pre-allocated chunks using a
@@ -79,7 +81,8 @@
  * <h3>How it works</h3>
  *
  * <ul>
- *   <li>Allocates large chunks (default 4MB) from native memory via {@code Unsafe}.
+ *   <li>Allocates large chunks (default 4MB) from Netty direct buffers ({@code
+ *       UnpooledByteBufAllocator}).
  *   <li>For each small allocation request, bumps a pointer within the current active chunk.
  *   <li>Only switches to a new chunk when the current one has no room.
  *   <li>Reference-counts each chunk: when all sub-allocations are released, the chunk is recycled
@@ -106,7 +109,7 @@
  * }</pre>
  *
  * <p>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<Chunk> freeChunks = new ArrayDeque<>();
 
+        /** Total direct memory, in bytes, currently allocated by this factory. */
+        private final AtomicLong directMemoryUsedBytes = 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);
+                directMemoryUsedBytes.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);
+            directMemoryUsedBytes.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 used by this factory. */
+        public long getDirectMemoryUsedBytes() {
+            return directMemoryUsedBytes.get();
+        }
+
+        private void destroyChunk(Chunk chunk) {
+            long cap = chunk.capacity;
+            chunk.destroy();
+            directMemoryUsedBytes.addAndGet(-cap);
+        }
+
+        void decrementDirectMemoryBytes(long size) {
+            directMemoryUsedBytes.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..ad6d8d79f64 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 testDirectMemoryUsedBytesTracksChunksAndDirectAllocations() {
+        assertThat(factory.getDirectMemoryUsedBytes()).isZero();
+
+        try (ArrowBuf small = allocator.buffer(64)) {
+            assertThat(small.capacity()).isEqualTo(64L);
+            assertThat(factory.getDirectMemoryUsedBytes()).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.getDirectMemoryUsedBytes())
+                        .isEqualTo(TEST_CHUNK_SIZE + direct.capacity());
+            }
+
+            assertThat(factory.getDirectMemoryUsedBytes()).isEqualTo(TEST_CHUNK_SIZE);
+
+            try (ArrowBuf anotherSmall = allocator.buffer(64)) {
+                assertThat(anotherSmall.capacity()).isEqualTo(64L);
+                // reuse the same chunk.
+                assertThat(factory.getDirectMemoryUsedBytes()).isEqualTo(TEST_CHUNK_SIZE);
+            }
+        }
+
+        // The empty active chunk remains cached until the factory is closed.
+        assertThat(factory.getDirectMemoryUsedBytes()).isEqualTo(TEST_CHUNK_SIZE);
+        factory.close();
+        assertThat(factory.getDirectMemoryUsedBytes()).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.getDirectMemoryUsedBytes()).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.getDirectMemoryUsedBytes()).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.getDirectMemoryUsedBytes()).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.getDirectMemoryUsedBytes()).isEqualTo(directCapacity);
+
+        // Close factory while the direct ArrowBuf is still alive.
+        localFactory.close();
+
+        // Memory is still tracked (ArrowBuf not yet released).
+        assertThat(localFactory.getDirectMemoryUsedBytes()).isEqualTo(directCapacity);
+
+        // Release the ArrowBuf — should properly free the direct ByteBuf and decrement.
+        directBuf.close();
+
+        // The direct allocation is freed.
+        assertThat(localFactory.getDirectMemoryUsedBytes()).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.getDirectMemoryUsedBytes()).isEqualTo(directCapacity);
+            directBuf.close();
+            assertThat(localFactory.getDirectMemoryUsedBytes()).isZero();
+        }
+
+        localFactory.close();
+        assertThat(localFactory.getDirectMemoryUsedBytes()).isZero();
         localAllocator.close();
     }
 }

From 0f368c31ee363fdafedb279b4df23799b68b943b Mon Sep 17 00:00:00 2001
From: Hongshun Wang <loserwang1024@gmail>
Date: Tue, 22 Sep 2026 10:44:21 +0800
Subject: [PATCH 2/2] Modified based on CR

Co-Authored-By: Qoder <noreply@qoder.com>
Co-Authored-By: Codex <noreply@openai.com>
AI-Model: gpt-6-astra
AI-Contributed/Feature: 24/24
AI-Contributed/UT: 46/46
---
 .../fluss/client/write/RecordAccumulator.java |  4 +--
 .../client/write/RecordAccumulatorTest.java   | 12 ++++---
 .../org/apache/fluss/metrics/MetricNames.java |  4 +--
 .../memory/ChunkedAllocationManager.java      | 16 ++++-----
 .../memory/ChunkedAllocationManagerTest.java  | 34 +++++++++----------
 5 files changed, 36 insertions(+), 34 deletions(-)

diff --git a/fluss-client/src/main/java/org/apache/fluss/client/write/RecordAccumulator.java b/fluss-client/src/main/java/org/apache/fluss/client/write/RecordAccumulator.java
index 5ff98c64406..45f45d89e9d 100644
--- a/fluss-client/src/main/java/org/apache/fluss/client/write/RecordAccumulator.java
+++ b/fluss-client/src/main/java/org/apache/fluss/client/write/RecordAccumulator.java
@@ -213,8 +213,8 @@ private void registerMetrics(WriterMetricGroup writerMetricGroup) {
                 MetricNames.WRITER_ACCUMULATOR_ARROW_MEMORY_USED_BYTES,
                 bufferAllocator::getAllocatedMemory);
         writerMetricGroup.gauge(
-                MetricNames.WRITER_ACCUMULATOR_DIRECT_MEMORY_USED_BYTES,
-                chunkedFactory::getDirectMemoryUsedBytes);
+                MetricNames.WRITER_ACCUMULATOR_DIRECT_MEMORY_ALLOCATED_BYTES,
+                chunkedFactory::getDirectMemoryAllocatedBytes);
     }
 
     /** Assigns and appends a record using the layout owned by its write context. */
diff --git a/fluss-client/src/test/java/org/apache/fluss/client/write/RecordAccumulatorTest.java b/fluss-client/src/test/java/org/apache/fluss/client/write/RecordAccumulatorTest.java
index f2d2828bd53..50f42464c86 100644
--- a/fluss-client/src/test/java/org/apache/fluss/client/write/RecordAccumulatorTest.java
+++ b/fluss-client/src/test/java/org/apache/fluss/client/write/RecordAccumulatorTest.java
@@ -163,7 +163,7 @@ void testAccumulatorMemoryMetrics() throws Exception {
                     .containsKeys(
                             MetricNames.WRITER_ACCUMULATOR_HEAP_MEMORY_USED_BYTES,
                             MetricNames.WRITER_ACCUMULATOR_ARROW_MEMORY_USED_BYTES,
-                            MetricNames.WRITER_ACCUMULATOR_DIRECT_MEMORY_USED_BYTES);
+                            MetricNames.WRITER_ACCUMULATOR_DIRECT_MEMORY_ALLOCATED_BYTES);
 
             Gauge<?> heapMemoryUsed =
                     (Gauge<?>)
@@ -173,14 +173,16 @@ void testAccumulatorMemoryMetrics() throws Exception {
                     (Gauge<?>)
                             metrics.getMetrics()
                                     .get(MetricNames.WRITER_ACCUMULATOR_ARROW_MEMORY_USED_BYTES);
-            Gauge<?> directMemoryUsedBytes =
+            Gauge<?> directMemoryAllocatedBytes =
                     (Gauge<?>)
                             metrics.getMetrics()
-                                    .get(MetricNames.WRITER_ACCUMULATOR_DIRECT_MEMORY_USED_BYTES);
+                                    .get(
+                                            MetricNames
+                                                    .WRITER_ACCUMULATOR_DIRECT_MEMORY_ALLOCATED_BYTES);
 
             assertThat(((Number) heapMemoryUsed.getValue()).longValue()).isZero();
             assertThat(((Number) arrowMemoryUsed.getValue()).longValue()).isZero();
-            assertThat(((Number) directMemoryUsedBytes.getValue()).longValue()).isZero();
+            assertThat(((Number) directMemoryAllocatedBytes.getValue()).longValue()).isZero();
 
             bucketAssigner.setBucketId(0);
             accum.append(
@@ -191,7 +193,7 @@ DATA1_TABLE_INFO, DATA1_PHYSICAL_TABLE_PATH, row(1, "a"), null),
 
             long heapMemory = ((Number) heapMemoryUsed.getValue()).longValue();
             long arrowMemory = ((Number) arrowMemoryUsed.getValue()).longValue();
-            long directMemory = ((Number) directMemoryUsedBytes.getValue()).longValue();
+            long directMemory = ((Number) directMemoryAllocatedBytes.getValue()).longValue();
             assertThat(heapMemory).isPositive();
             assertThat(arrowMemory).isPositive();
             assertThat(directMemory).isGreaterThanOrEqualTo(arrowMemory);
diff --git a/fluss-common/src/main/java/org/apache/fluss/metrics/MetricNames.java b/fluss-common/src/main/java/org/apache/fluss/metrics/MetricNames.java
index dd9c844efff..496a212e48b 100644
--- a/fluss-common/src/main/java/org/apache/fluss/metrics/MetricNames.java
+++ b/fluss-common/src/main/java/org/apache/fluss/metrics/MetricNames.java
@@ -342,8 +342,8 @@ public class MetricNames {
             "accumulatorHeapMemoryUsedBytes";
     public static final String WRITER_ACCUMULATOR_ARROW_MEMORY_USED_BYTES =
             "accumulatorArrowMemoryUsedBytes";
-    public static final String WRITER_ACCUMULATOR_DIRECT_MEMORY_USED_BYTES =
-            "accumulatorDirectMemoryUsedBytes";
+    public static final String WRITER_ACCUMULATOR_DIRECT_MEMORY_ALLOCATED_BYTES =
+            "accumulatorDirectMemoryAllocatedBytes";
 
     // for scanner
     public static final String SCANNER_TIME_MS_BETWEEN_POLL = "timeMsBetweenPoll";
diff --git a/fluss-common/src/main/java/org/apache/fluss/shaded/arrow/org/apache/arrow/memory/ChunkedAllocationManager.java b/fluss-common/src/main/java/org/apache/fluss/shaded/arrow/org/apache/arrow/memory/ChunkedAllocationManager.java
index b032c657bdb..cc14340f7f8 100644
--- a/fluss-common/src/main/java/org/apache/fluss/shaded/arrow/org/apache/arrow/memory/ChunkedAllocationManager.java
+++ b/fluss-common/src/main/java/org/apache/fluss/shaded/arrow/org/apache/arrow/memory/ChunkedAllocationManager.java
@@ -309,7 +309,7 @@ public static class ChunkedFactory implements AllocationManager.Factory {
         private final Deque<Chunk> freeChunks = new ArrayDeque<>();
 
         /** Total direct memory, in bytes, currently allocated by this factory. */
-        private final AtomicLong directMemoryUsedBytes = new AtomicLong();
+        private final AtomicLong directMemoryAllocatedBytes = new AtomicLong();
 
         /** Set to true when {@link #close()} is called. */
         private boolean closed;
@@ -343,7 +343,7 @@ public synchronized AllocationManager create(
                             "Allocation size " + size + " exceeds maximum " + Integer.MAX_VALUE);
                 }
                 ByteBuf directByteBuf = UnpooledByteBufAllocator.DEFAULT.directBuffer((int) size);
-                directMemoryUsedBytes.addAndGet(size);
+                directMemoryAllocatedBytes.addAndGet(size);
                 return new ChunkedAllocationManager(accountingAllocator, directByteBuf, size, this);
             }
 
@@ -375,7 +375,7 @@ private synchronized Chunk obtainChunk() {
                 return recycled;
             }
             Chunk chunk = new Chunk(chunkSize, this);
-            directMemoryUsedBytes.addAndGet(chunkSize);
+            directMemoryAllocatedBytes.addAndGet(chunkSize);
             return chunk;
         }
 
@@ -439,19 +439,19 @@ public synchronized void close() {
             }
         }
 
-        /** Returns the direct memory, in bytes, currently used by this factory. */
-        public long getDirectMemoryUsedBytes() {
-            return directMemoryUsedBytes.get();
+        /** 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();
-            directMemoryUsedBytes.addAndGet(-cap);
+            directMemoryAllocatedBytes.addAndGet(-cap);
         }
 
         void decrementDirectMemoryBytes(long size) {
-            directMemoryUsedBytes.addAndGet(-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 ad6d8d79f64..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
@@ -72,33 +72,33 @@ void testSmallAllocationPackedIntoSameChunk() {
     }
 
     @Test
-    void testDirectMemoryUsedBytesTracksChunksAndDirectAllocations() {
-        assertThat(factory.getDirectMemoryUsedBytes()).isZero();
+    void testDirectMemoryAllocatedBytesTracksChunksAndDirectAllocations() {
+        assertThat(factory.getDirectMemoryAllocatedBytes()).isZero();
 
         try (ArrowBuf small = allocator.buffer(64)) {
             assertThat(small.capacity()).isEqualTo(64L);
-            assertThat(factory.getDirectMemoryUsedBytes()).isEqualTo(TEST_CHUNK_SIZE);
+            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.getDirectMemoryUsedBytes())
+                assertThat(factory.getDirectMemoryAllocatedBytes())
                         .isEqualTo(TEST_CHUNK_SIZE + direct.capacity());
             }
 
-            assertThat(factory.getDirectMemoryUsedBytes()).isEqualTo(TEST_CHUNK_SIZE);
+            assertThat(factory.getDirectMemoryAllocatedBytes()).isEqualTo(TEST_CHUNK_SIZE);
 
             try (ArrowBuf anotherSmall = allocator.buffer(64)) {
                 assertThat(anotherSmall.capacity()).isEqualTo(64L);
                 // reuse the same chunk.
-                assertThat(factory.getDirectMemoryUsedBytes()).isEqualTo(TEST_CHUNK_SIZE);
+                assertThat(factory.getDirectMemoryAllocatedBytes()).isEqualTo(TEST_CHUNK_SIZE);
             }
         }
 
         // The empty active chunk remains cached until the factory is closed.
-        assertThat(factory.getDirectMemoryUsedBytes()).isEqualTo(TEST_CHUNK_SIZE);
+        assertThat(factory.getDirectMemoryAllocatedBytes()).isEqualTo(TEST_CHUNK_SIZE);
         factory.close();
-        assertThat(factory.getDirectMemoryUsedBytes()).isZero();
+        assertThat(factory.getDirectMemoryAllocatedBytes()).isZero();
     }
 
     @Test
@@ -320,7 +320,7 @@ void testCloseWhileArrowBufsAliveDoesNotLeakMemory() {
         for (int i = 0; i < 4; i++) {
             bufs.add(localAllocator.buffer(64));
         }
-        assertThat(localFactory.getDirectMemoryUsedBytes()).isEqualTo(TEST_CHUNK_SIZE);
+        assertThat(localFactory.getDirectMemoryAllocatedBytes()).isEqualTo(TEST_CHUNK_SIZE);
 
         // Close factory while ArrowBufs are still alive.
 
@@ -328,14 +328,14 @@ void testCloseWhileArrowBufsAliveDoesNotLeakMemory() {
 
         // Active chunk was not destroyed (subAllocCount > 0), just nulled.
         assertThat(localFactory.getActiveChunk()).isNull();
-        assertThat(localFactory.getDirectMemoryUsedBytes()).isEqualTo(TEST_CHUNK_SIZE);
+        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.getDirectMemoryUsedBytes()).isZero();
+        assertThat(localFactory.getDirectMemoryAllocatedBytes()).isZero();
         localAllocator.close();
     }
 
@@ -349,19 +349,19 @@ void testDirectAllocationReleaseAfterFactoryClose() {
         // Allocate a large buffer that bypasses chunk sub-allocation.
         ArrowBuf directBuf = localAllocator.buffer(TEST_CHUNK_SIZE + 1);
         long directCapacity = directBuf.capacity();
-        assertThat(localFactory.getDirectMemoryUsedBytes()).isEqualTo(directCapacity);
+        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.getDirectMemoryUsedBytes()).isEqualTo(directCapacity);
+        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.getDirectMemoryUsedBytes()).isZero();
+        assertThat(localFactory.getDirectMemoryAllocatedBytes()).isZero();
         localAllocator.close();
     }
 
@@ -389,13 +389,13 @@ void testDirectAllocationMultipleCycles() {
         for (int cycle = 0; cycle < 5; cycle++) {
             ArrowBuf directBuf = localAllocator.buffer(directSize);
             long directCapacity = directBuf.capacity();
-            assertThat(localFactory.getDirectMemoryUsedBytes()).isEqualTo(directCapacity);
+            assertThat(localFactory.getDirectMemoryAllocatedBytes()).isEqualTo(directCapacity);
             directBuf.close();
-            assertThat(localFactory.getDirectMemoryUsedBytes()).isZero();
+            assertThat(localFactory.getDirectMemoryAllocatedBytes()).isZero();
         }
 
         localFactory.close();
-        assertThat(localFactory.getDirectMemoryUsedBytes()).isZero();
+        assertThat(localFactory.getDirectMemoryAllocatedBytes()).isZero();
         localAllocator.close();
     }
 }