From e87fc5cfdc3be417086f22409379a4a27d0ea318 Mon Sep 17 00:00:00 2001 From: Hongze Zhang Date: Tue, 30 Nov 2021 15:09:44 +0800 Subject: [PATCH] AllocationListener should be called in BufferLedger#transferBalance(...) (#51) --- .../main/java/org/apache/arrow/memory/BufferLedger.java | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/java/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferLedger.java b/java/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferLedger.java index 472897a64058e..346c4b30bc9f7 100644 --- a/java/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferLedger.java +++ b/java/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferLedger.java @@ -369,8 +369,12 @@ boolean transferBalance(final ReferenceManager targetReferenceManager) { targetReferenceManager.getAllocator().getName()); } - boolean overlimit = targetAllocator.forceAllocate(memoryChunkManager.getSize()); - allocator.releaseBytes(memoryChunkManager.getSize()); + long size = memoryChunkManager.getSize(); + targetAllocator.getListener().onPreAllocation(size); + boolean overlimit = targetAllocator.forceAllocate(size); + targetAllocator.getListener().onAllocation(size); + allocator.releaseBytes(size); + allocator.getListener().onRelease(size); // since the transfer can only happen from the owning reference manager, // we need to set the target ref manager as the new owning ref manager // for the chunk of memory in MemoryChunkManager