[GitHub] [arrow] davisusanibar commented on a change in pull request #11800: ARROW-14923: [Java] AllocationListener should be called during ownership transferring

[GitBox <gi...@apache.org>]

davisusanibar commented on a change in pull request #11800:
URL: https://github.com/apache/arrow/pull/11800#discussion_r760802379



##########
File path: java/memory/memory-netty/src/test/java/org/apache/arrow/memory/TestBaseAllocator.java
##########
@@ -159,6 +159,87 @@ public void testAllocator_transferOwnership() throws Exception {
     }
   }
 
+  @Test
+  public void testAllocator_transferOwnershipWithDifferentListeners() throws Exception {
+    try (final RootAllocator rootAllocator =
+             new RootAllocator(MAX_ALLOCATION)) {
+      TestAllocationListener l1 = new TestAllocationListener();
+      TestAllocationListener l2 = new TestAllocationListener();
+      final BufferAllocator childAllocator1 =
+          rootAllocator.newChildAllocator("changeOwnership1", l1, 0, MAX_ALLOCATION);
+      final BufferAllocator childAllocator2 =
+          rootAllocator.newChildAllocator("changeOwnership2", l2, 0, MAX_ALLOCATION);
+
+      ArrowBuf buffer1 = childAllocator1.buffer(1024L);
+
+      // initial states
+      assertEquals(1, l1.getNumCalls());
+      assertEquals(0, l1.getNumReleaseCalls());
+      assertEquals(0, l2.getNumCalls());
+      assertEquals(0, l2.getNumReleaseCalls());
+
+      // do transfer
+      OwnershipTransferResult transferResult = buffer1.getReferenceManager()
+          .transferOwnership(buffer1, childAllocator2);
+      assertTrue(transferResult.getAllocationFit());
+      ArrowBuf buffer2 = transferResult.getTransferredBuffer();
+
+      // reservations from l1 should also be transferred to l2 during buffer transferring
+      assertEquals(1, l1.getNumCalls());
+      assertEquals(1, l1.getNumReleaseCalls());

Review comment:
       Please add a comment to mark that this assertion should failed without this change

##########
File path: java/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferLedger.java
##########
@@ -364,8 +364,12 @@ boolean transferBalance(final ReferenceManager targetReferenceManager) {
             targetReferenceManager.getAllocator().getName());
       }
 
-      boolean overlimit = targetAllocator.forceAllocate(allocationManager.getSize());
-      allocator.releaseBytes(allocationManager.getSize());
+      long size = allocationManager.getSize();
+      targetAllocator.getListener().onPreAllocation(size);
+      boolean overlimit = targetAllocator.forceAllocate(size);
+      targetAllocator.getListener().onAllocation(size);
+      allocator.releaseBytes(size);
+      allocator.getListener().onRelease(size);

Review comment:
       If we do not validate "assertEquals(1, l1.getNumReleaseCalls())" or "assertEquals(1, l2.getNumCalls())", question: What could be the errors in current programs to this changes should be needed ?




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org