[GitHub] [arrow] lidavidm commented on a diff in pull request #34231: GH-34230: [Java] Call allocation listener on BaseAllocator#wrapForeignAllocation

["lidavidm (via GitHub)" <gi...@apache.org>]

lidavidm commented on code in PR #34231:
URL: https://github.com/apache/arrow/pull/34231#discussion_r1109781042


##########
java/memory/memory-netty/src/test/java/org/apache/arrow/memory/CountingAllocationListener.java:
##########
@@ -0,0 +1,111 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.arrow.memory;
+
+// Allocation listener

Review Comment:
   nit, but please stick to JavaDoc syntax for this?



##########
java/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferAllocator.java:
##########
@@ -248,22 +248,5 @@ default RoundingPolicy getRoundingPolicy() {
    *
    * @param allocation The underlying allocation.
    */
-  default ArrowBuf wrapForeignAllocation(ForeignAllocation allocation) {
-    try {
-      forceAllocate(allocation.getSize());
-      final AllocationManager manager = new ForeignAllocationManager(this, allocation);
-      final BufferLedger ledger = manager.associate(this);
-      final ArrowBuf buf =
-          new ArrowBuf(ledger, /*bufferManager=*/null, allocation.getSize(), allocation.memoryAddress());
-      buf.writerIndex(allocation.getSize());
-      return buf;
-    } catch (Throwable t) {
-      try {
-        allocation.release0();
-      } catch (Throwable e) {
-        t.addSuppressed(e);
-      }
-      throw t;
-    }
-  }
+  ArrowBuf wrapForeignAllocation(ForeignAllocation allocation);

Review Comment:
   This is a breaking change. Can we at least have a default implementation that throws UnsupportedOperationException?



##########
java/memory/memory-netty/src/test/java/org/apache/arrow/memory/TestForeignAllocation.java:
##########
@@ -54,6 +54,55 @@ public void wrapForeignAllocation() {
     assertEquals(0, allocator.getAllocatedMemory());
   }
 
+  @Test
+  public void wrapForeignAllocationWithAllocationListener() {
+    final long bufferSize = 16;
+
+    final CountingAllocationListener listener = new CountingAllocationListener();
+    try (BufferAllocator listenedAllocator =
+        allocator.newChildAllocator("child", listener, 0L, allocator.getLimit())) {
+      UnsafeForeignAllocation allocation = new UnsafeForeignAllocation(bufferSize);
+      try {
+        assertEquals(0, listenedAllocator.getAllocatedMemory());
+        ArrowBuf buf = listenedAllocator.wrapForeignAllocation(allocation);
+        assertEquals(bufferSize, buf.capacity());
+        assertEquals(16, listener.getCurrentMem());
+        buf.close();
+        assertEquals(0, listener.getCurrentMem());
+        assertTrue(allocation.released);
+      } finally {
+        allocation.release0();
+      }
+      assertEquals(0, listenedAllocator.getAllocatedMemory());
+    }
+    assertEquals(1, listener.getNumPreCalls());
+    assertEquals(1, listener.getNumCalls());
+    assertEquals(1, listener.getNumReleaseCalls());
+    assertEquals(16, listener.getTotalMem());
+  }
+
+  @Test(expected = OutOfMemoryException.class)
+  public void wrapForeignAllocationFailedWithAllocationListener() {

Review Comment:
   Can we also test the case where the listener frees some memory and the allocation succeeds?



##########
java/memory/memory-core/src/main/java/org/apache/arrow/memory/BaseAllocator.java:
##########
@@ -233,6 +233,41 @@ private void childClosed(final BaseAllocator childAllocator) {
     listener.onChildRemoved(this, childAllocator);
   }
 
+  @Override
+  public ArrowBuf wrapForeignAllocation(ForeignAllocation allocation) {
+    assertOpen();
+    final long size = allocation.getSize();
+    listener.onPreAllocation(size);
+    AllocationOutcome outcome = this.allocateBytes(size);

Review Comment:
   Hmm, if we don't force allocation the behavior is different - the memory is already allocated, we're just trying to account for it. This will 'eagerly' release resources (via the listener) to make room while the existing behavior is to do it 'lazily' (the next allocation will trigger things if needed). Is that what you want?



##########
java/memory/memory-core/src/main/java/org/apache/arrow/memory/BaseAllocator.java:
##########
@@ -233,6 +233,41 @@ private void childClosed(final BaseAllocator childAllocator) {
     listener.onChildRemoved(this, childAllocator);
   }
 
+  @Override
+  public ArrowBuf wrapForeignAllocation(ForeignAllocation allocation) {
+    assertOpen();
+    final long size = allocation.getSize();
+    listener.onPreAllocation(size);
+    AllocationOutcome outcome = this.allocateBytes(size);
+    if (!outcome.isOk()) {
+      if (listener.onFailedAllocation(size, outcome)) {
+        // Second try, in case the listener can do something about it
+        outcome = this.allocateBytes(size);
+      }
+      if (!outcome.isOk()) {
+        throw new OutOfMemoryException(createErrorMsg(this, size,
+            size), outcome.getDetails());
+      }
+    }
+    try {
+      final AllocationManager manager = new ForeignAllocationManager(this, allocation);
+      final BufferLedger ledger = manager.associate(this);
+      final ArrowBuf buf =
+          new ArrowBuf(ledger, /*bufferManager=*/null, size, allocation.memoryAddress());
+      buf.writerIndex(size);
+      listener.onAllocation(size);
+      return buf;
+    } catch (Throwable t) {
+      try {
+        releaseBytes(size);
+        allocation.release0();

Review Comment:
   can this be in a separate try-catch?



-- 
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