[arrow] branch main updated: GH-35960: [Java] Detect overflow in allocation (#36185)

[li...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

lidavidm pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow.git


The following commit(s) were added to refs/heads/main by this push:
     new 1ab00aefb4 GH-35960: [Java] Detect overflow in allocation (#36185)
1ab00aefb4 is described below

commit 1ab00aefb493bb4f38bde6e41c92e2dfe1782452
Author: David Li <li...@gmail.com>
AuthorDate: Thu Jun 29 08:15:04 2023 -0400

    GH-35960: [Java] Detect overflow in allocation (#36185)
    
    
    
    ### Rationale for this change
    
    The Java allocator wasn't detecting overflow (which is admittedly a corner case since you can't actually allocate that many bytes), causing an unexpected exception.
    
    ### What changes are included in this PR?
    
    Detect overflow and provide the right exception.
    
    ### Are these changes tested?
    
    Yes
    
    ### Are there any user-facing changes?
    
    No
    * Closes: #35960
    
    Authored-by: David Li <li...@gmail.com>
    Signed-off-by: David Li <li...@gmail.com>
---
 .../java/org/apache/arrow/memory/Accountant.java   |  9 ++++--
 .../org/apache/arrow/memory/TestBaseAllocator.java | 33 ++++++++++++++++++++++
 2 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/java/memory/memory-core/src/main/java/org/apache/arrow/memory/Accountant.java b/java/memory/memory-core/src/main/java/org/apache/arrow/memory/Accountant.java
index 42dac7b8c6..87769dd122 100644
--- a/java/memory/memory-core/src/main/java/org/apache/arrow/memory/Accountant.java
+++ b/java/memory/memory-core/src/main/java/org/apache/arrow/memory/Accountant.java
@@ -169,9 +169,14 @@ class Accountant implements AutoCloseable {
    */
   private AllocationOutcome.Status allocate(final long size, final boolean incomingUpdatePeak,
       final boolean forceAllocation, AllocationOutcomeDetails details) {
-    final long newLocal = locallyHeldMemory.addAndGet(size);
+    final long oldLocal = locallyHeldMemory.getAndAdd(size);
+    final long newLocal = oldLocal + size;
+    // Borrowed from Math.addExact (but avoid exception here)
+    // Overflow if result has opposite sign of both arguments
+    // No need to reset locallyHeldMemory on overflow; allocateBytesInternal will releaseBytes on failure
+    final boolean overflow = ((oldLocal ^ newLocal) & (size ^ newLocal)) < 0;
     final long beyondReservation = newLocal - reservation;
-    final boolean beyondLimit = newLocal > allocationLimit.get();
+    final boolean beyondLimit = overflow || newLocal > allocationLimit.get();
     final boolean updatePeak = forceAllocation || (incomingUpdatePeak && !beyondLimit);
 
     if (details != null) {
diff --git a/java/memory/memory-netty/src/test/java/org/apache/arrow/memory/TestBaseAllocator.java b/java/memory/memory-netty/src/test/java/org/apache/arrow/memory/TestBaseAllocator.java
index 7c0df0e98e..7613d073f8 100644
--- a/java/memory/memory-netty/src/test/java/org/apache/arrow/memory/TestBaseAllocator.java
+++ b/java/memory/memory-netty/src/test/java/org/apache/arrow/memory/TestBaseAllocator.java
@@ -1134,6 +1134,39 @@ public class TestBaseAllocator {
     }
   }
 
+  @Test
+  public void testOverlimit() {
+    try (BufferAllocator allocator = new RootAllocator(1024)) {
+      try (BufferAllocator child1 = allocator.newChildAllocator("ChildA", 0, 1024);
+           BufferAllocator child2 = allocator.newChildAllocator("ChildB", 1024, 1024)) {
+        assertThrows(OutOfMemoryException.class, () -> {
+          ArrowBuf buf1 = child1.buffer(8);
+          buf1.close();
+        });
+        assertEquals(0, child1.getAllocatedMemory());
+        assertEquals(0, child2.getAllocatedMemory());
+        assertEquals(1024, allocator.getAllocatedMemory());
+      }
+    }
+  }
+
+  @Test
+  public void testOverlimitOverflow() {
+    // Regression test for https://github.com/apache/arrow/issues/35960
+    try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE)) {
+      try (BufferAllocator child1 = allocator.newChildAllocator("ChildA", 0, Long.MAX_VALUE);
+           BufferAllocator child2 = allocator.newChildAllocator("ChildB", Long.MAX_VALUE, Long.MAX_VALUE)) {
+        assertThrows(OutOfMemoryException.class, () -> {
+          ArrowBuf buf1 = child1.buffer(1024);
+          buf1.close();
+        });
+        assertEquals(0, child1.getAllocatedMemory());
+        assertEquals(0, child2.getAllocatedMemory());
+        assertEquals(Long.MAX_VALUE, allocator.getAllocatedMemory());
+      }
+    }
+  }
+
   public void assertEquiv(ArrowBuf origBuf, ArrowBuf newBuf) {
     assertEquals(origBuf.readerIndex(), newBuf.readerIndex());
     assertEquals(origBuf.writerIndex(), newBuf.writerIndex());