Repository: arrow Updated Branches: refs/heads/master 292618327 -> f67974b19
ARROW-700: Add headroom interface for allocator Author: Julien Le Dem <[email protected]> Closes #424 from julienledem/headroom and squashes the following commits: 2aab160 [Julien Le Dem] ARROW-700: Add headroom interface for allocator Project: http://git-wip-us.apache.org/repos/asf/arrow/repo Commit: http://git-wip-us.apache.org/repos/asf/arrow/commit/f67974b1 Tree: http://git-wip-us.apache.org/repos/asf/arrow/tree/f67974b1 Diff: http://git-wip-us.apache.org/repos/asf/arrow/diff/f67974b1 Branch: refs/heads/master Commit: f67974b190349c781509d2b1657331935f458f9b Parents: 2926183 Author: Julien Le Dem <[email protected]> Authored: Wed Mar 22 23:10:01 2017 -0400 Committer: Wes McKinney <[email protected]> Committed: Wed Mar 22 23:10:01 2017 -0400 ---------------------------------------------------------------------- .../main/java/org/apache/arrow/memory/Accountant.java | 14 ++++++++++++-- .../java/org/apache/arrow/memory/BufferAllocator.java | 8 ++++++++ .../java/org/apache/arrow/memory/TestAccountant.java | 13 +++++++++++-- 3 files changed, 31 insertions(+), 4 deletions(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/arrow/blob/f67974b1/java/memory/src/main/java/org/apache/arrow/memory/Accountant.java ---------------------------------------------------------------------- diff --git a/java/memory/src/main/java/org/apache/arrow/memory/Accountant.java b/java/memory/src/main/java/org/apache/arrow/memory/Accountant.java index 6ddc8f7..89329b2 100644 --- a/java/memory/src/main/java/org/apache/arrow/memory/Accountant.java +++ b/java/memory/src/main/java/org/apache/arrow/memory/Accountant.java @@ -18,12 +18,12 @@ package org.apache.arrow.memory; -import com.google.common.base.Preconditions; - import java.util.concurrent.atomic.AtomicLong; import javax.annotation.concurrent.ThreadSafe; +import com.google.common.base.Preconditions; + /** * Provides a concurrent way to manage account for memory usage without locking. Used as basis * for Allocators. All @@ -202,6 +202,7 @@ class Accountant implements AutoCloseable { /** * Close this Accountant. This will release any reservation bytes back to a parent Accountant. */ + @Override public void close() { // return memory reservation to parent allocator. if (parent != null) { @@ -248,6 +249,15 @@ class Accountant implements AutoCloseable { return peakAllocation.get(); } + public long getHeadroom(){ + long localHeadroom = allocationLimit.get() - locallyHeldMemory.get(); + if(parent == null){ + return localHeadroom; + } + + return Math.min(localHeadroom, parent.getHeadroom()); + } + /** * Describes the type of outcome that occurred when trying to account for allocation of memory. */ http://git-wip-us.apache.org/repos/asf/arrow/blob/f67974b1/java/memory/src/main/java/org/apache/arrow/memory/BufferAllocator.java ---------------------------------------------------------------------- diff --git a/java/memory/src/main/java/org/apache/arrow/memory/BufferAllocator.java b/java/memory/src/main/java/org/apache/arrow/memory/BufferAllocator.java index 81ffb1b..c05e9ac 100644 --- a/java/memory/src/main/java/org/apache/arrow/memory/BufferAllocator.java +++ b/java/memory/src/main/java/org/apache/arrow/memory/BufferAllocator.java @@ -106,6 +106,14 @@ public interface BufferAllocator extends AutoCloseable { public long getPeakMemoryAllocation(); /** + * Returns the amount of memory that can probably be allocated at this moment + * without exceeding this or any parents allocation maximum. + * + * @return Headroom in bytes + */ + public long getHeadroom(); + + /** * Create an allocation reservation. A reservation is a way of building up * a request for a buffer whose size is not known in advance. See * {@see AllocationReservation}. http://git-wip-us.apache.org/repos/asf/arrow/blob/f67974b1/java/memory/src/test/java/org/apache/arrow/memory/TestAccountant.java ---------------------------------------------------------------------- diff --git a/java/memory/src/test/java/org/apache/arrow/memory/TestAccountant.java b/java/memory/src/test/java/org/apache/arrow/memory/TestAccountant.java index 86bccf5..2624a4a 100644 --- a/java/memory/src/test/java/org/apache/arrow/memory/TestAccountant.java +++ b/java/memory/src/test/java/org/apache/arrow/memory/TestAccountant.java @@ -19,7 +19,6 @@ package org.apache.arrow.memory; import static org.junit.Assert.assertEquals; -import org.apache.arrow.memory.Accountant; import org.apache.arrow.memory.Accountant.AllocationOutcome; import org.junit.Assert; import org.junit.Test; @@ -36,6 +35,7 @@ public class TestAccountant { final Accountant parent = new Accountant(null, 0, Long.MAX_VALUE); ensureAccurateReservations(parent); assertEquals(0, parent.getAllocatedMemory()); + assertEquals(parent.getLimit() - parent.getAllocatedMemory(), parent.getHeadroom()); } @Test @@ -71,6 +71,7 @@ public class TestAccountant { } assertEquals(0, parent.getAllocatedMemory()); + assertEquals(parent.getLimit() - parent.getAllocatedMemory(), parent.getHeadroom()); } private void ensureAccurateReservations(Accountant outsideParent) { @@ -121,6 +122,9 @@ public class TestAccountant { // went beyond reservation, now in parent accountant assertEquals(3, parent.getAllocatedMemory()); + assertEquals(7, child.getHeadroom()); + assertEquals(7, parent.getHeadroom()); + { AllocationOutcome first = child.allocateBytes(7); assertEquals(AllocationOutcome.SUCCESS, first); @@ -135,9 +139,11 @@ public class TestAccountant { child.releaseBytes(9); assertEquals(1, child.getAllocatedMemory()); + assertEquals(8, child.getHeadroom()); // back to reservation size assertEquals(2, parent.getAllocatedMemory()); + assertEquals(8, parent.getHeadroom()); AllocationOutcome first = child.allocateBytes(10); assertEquals(AllocationOutcome.FAILED_PARENT, first); @@ -152,11 +158,14 @@ public class TestAccountant { // at new limit assertEquals(child.getAllocatedMemory(), 11); assertEquals(parent.getAllocatedMemory(), 11); - + assertEquals(-1, child.getHeadroom()); + assertEquals(-1, parent.getHeadroom()); child.releaseBytes(11); assertEquals(child.getAllocatedMemory(), 0); assertEquals(parent.getAllocatedMemory(), 2); + assertEquals(8, child.getHeadroom()); + assertEquals(8, parent.getHeadroom()); child.close(); parent.close();
