AllocationListener should be called in BufferLedger#transferBalance(.…

…..) (apache#51)
rui-mo · Apr 27, 2022 · 5859cc9 · 5859cc9
1 parent d891503
commit 5859cc9
Showing 1 changed file with 6 additions and 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
@@ -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