Remove no-fail memory reservation parameter

prestodb · Apr 15, 2015 · 51320bb · 51320bb
1 parent dd5984c
commit 51320bb
Show file tree

Hide file tree

Showing 6 changed files with 40 additions and 52 deletions.
diff --git a/presto-main/src/main/java/com/facebook/presto/operator/DriverContext.java b/presto-main/src/main/java/com/facebook/presto/operator/DriverContext.java
@@ -190,13 +190,10 @@ public DataSize getOperatorPreAllocatedMemory()
         return pipelineContext.getOperatorPreAllocatedMemory();
     }
 
-    public boolean reserveMemory(long bytes)
+    public void reserveMemory(long bytes)
     {
-        boolean result = pipelineContext.reserveMemory(bytes);
-        if (result) {
-            memoryReservation.getAndAdd(bytes);
-        }
-        return result;
+        pipelineContext.reserveMemory(bytes);
+        memoryReservation.getAndAdd(bytes);
     }
 
     public void freeMemory(long bytes)

diff --git a/presto-main/src/main/java/com/facebook/presto/operator/HashAggregationOperator.java b/presto-main/src/main/java/com/facebook/presto/operator/HashAggregationOperator.java
@@ -219,11 +219,6 @@ public Page getOutput()
                 return null;
             }
 
-            // Only partial aggregation can flush early. Also, check that we are not flushing tiny bits at a time
-            if (!finishing && step != Step.PARTIAL) {
-                throw new ExceededMemoryLimitException(operatorContext.getMaxMemorySize());
-            }
-
             outputIterator = aggregationBuilder.build();
             aggregationBuilder = null;
 
@@ -255,6 +250,7 @@ private static class GroupByHashAggregationBuilder
         private final GroupByHash groupByHash;
         private final List<Aggregator> aggregators;
         private final OperatorContext operatorContext;
+        private final boolean partial;
 
         private GroupByHashAggregationBuilder(
                 List<AccumulatorFactory> accumulatorFactories,
@@ -267,6 +263,7 @@ private GroupByHashAggregationBuilder(
         {
             this.groupByHash = createGroupByHash(groupByTypes, Ints.toArray(groupByChannels), hashChannel, expectedGroups);
             this.operatorContext = operatorContext;
+            this.partial = (step == Step.PARTIAL);
 
             // wrapper each function with an aggregator
             ImmutableList.Builder<Aggregator> builder = ImmutableList.builder();
@@ -302,7 +299,16 @@ public boolean isFull()
             if (memorySize < 0) {
                 memorySize = 0;
             }
-            return operatorContext.setMemoryReservation(memorySize, true) != memorySize;
+            try {
+                operatorContext.setMemoryReservation(memorySize);
+                return false;
+            }
+            catch (ExceededMemoryLimitException e) {
+                if (partial) {
+                    return true;
+                }
+                throw e;
+            }
         }
 
         public Iterator<Page> build()

diff --git a/presto-main/src/main/java/com/facebook/presto/operator/OperatorContext.java b/presto-main/src/main/java/com/facebook/presto/operator/OperatorContext.java
@@ -203,29 +203,14 @@ public DataSize getOperatorPreAllocatedMemory()
         return driverContext.getOperatorPreAllocatedMemory();
     }
 
-    public boolean reserveMemory(long bytes)
-    {
-        return reserveMemory(bytes, false);
-    }
-
-    public boolean reserveMemory(long bytes, boolean noFail)
+    public void reserveMemory(long bytes)
     {
+        driverContext.reserveMemory(bytes);
         long newReservation = memoryReservation.getAndAdd(bytes);
         if (newReservation > maxMemoryReservation) {
             memoryReservation.getAndAdd(-bytes);
-            if (!noFail) {
-                throw new ExceededMemoryLimitException(getMaxMemorySize());
-            }
-            return false;
+            throw new ExceededMemoryLimitException(getMaxMemorySize());
         }
-        boolean result = driverContext.reserveMemory(bytes);
-        if (!result) {
-            memoryReservation.getAndAdd(-bytes);
-            if (!noFail) {
-                throw new ExceededMemoryLimitException(getMaxMemorySize());
-            }
-        }
-        return result;
     }
 
     public void freeMemory(long bytes)
@@ -234,25 +219,18 @@ public void freeMemory(long bytes)
         memoryReservation.getAndAdd(-bytes);
     }
 
-    public synchronized long setMemoryReservation(long newMemoryReservation)
-    {
-        return setMemoryReservation(newMemoryReservation, false);
-    }
-
-    public synchronized long setMemoryReservation(long newMemoryReservation, boolean noFail)
+    public synchronized void setMemoryReservation(long newMemoryReservation)
     {
         checkArgument(newMemoryReservation >= 0, "newMemoryReservation is negative");
 
         long delta = newMemoryReservation - memoryReservation.get();
 
         if (delta > 0) {
-            reserveMemory(delta, noFail);
+            reserveMemory(delta);
         }
         else {
             freeMemory(-delta);
         }
-
-        return memoryReservation.get();
     }
 
     public void setInfoSupplier(Supplier<Object> infoSupplier)

diff --git a/presto-main/src/main/java/com/facebook/presto/operator/PipelineContext.java b/presto-main/src/main/java/com/facebook/presto/operator/PipelineContext.java
@@ -198,13 +198,10 @@ public DataSize getOperatorPreAllocatedMemory()
         return taskContext.getOperatorPreAllocatedMemory();
     }
 
-    public synchronized boolean reserveMemory(long bytes)
+    public synchronized void reserveMemory(long bytes)
     {
-        boolean result = taskContext.reserveMemory(bytes);
-        if (result) {
-            memoryReservation.getAndAdd(bytes);
-        }
-        return result;
+        taskContext.reserveMemory(bytes);
+        memoryReservation.getAndAdd(bytes);
     }
 
     public synchronized void freeMemory(long bytes)

diff --git a/presto-main/src/main/java/com/facebook/presto/operator/TaskContext.java b/presto-main/src/main/java/com/facebook/presto/operator/TaskContext.java
@@ -13,6 +13,7 @@
  */
 package com.facebook.presto.operator;
 
+import com.facebook.presto.ExceededMemoryLimitException;
 import com.facebook.presto.Session;
 import com.facebook.presto.execution.StateMachine.StateChangeListener;
 import com.facebook.presto.execution.TaskId;
@@ -180,15 +181,14 @@ public DataSize getOperatorPreAllocatedMemory()
         return operatorPreAllocatedMemory;
     }
 
-    public synchronized boolean reserveMemory(long bytes)
+    public synchronized void reserveMemory(long bytes)
     {
         checkArgument(bytes >= 0, "bytes is negative");
 
         if (memoryReservation.get() + bytes > maxMemory) {
-            return false;
+            throw new ExceededMemoryLimitException(getMaxMemorySize());
         }
         memoryReservation.getAndAdd(bytes);
-        return true;
     }
 
     public synchronized void freeMemory(long bytes)

diff --git a/presto-main/src/main/java/com/facebook/presto/operator/TopNOperator.java b/presto-main/src/main/java/com/facebook/presto/operator/TopNOperator.java
@@ -182,6 +182,7 @@ public void addInput(Page page)
         if (topNBuilder == null) {
             topNBuilder = new TopNBuilder(
                     n,
+                    partial,
                     sortTypes,
                     sortChannels,
                     sortOrders,
@@ -211,9 +212,6 @@ public Page getOutput()
                 outputIterator = topNBuilder.build();
                 topNBuilder = null;
             }
-            else {
-                throw new ExceededMemoryLimitException(operatorContext.getMaxMemorySize());
-            }
         }
 
         pageBuilder.reset();
@@ -232,6 +230,7 @@ public Page getOutput()
     private static class TopNBuilder
     {
         private final int n;
+        private final boolean partial;
         private final List<Type> sortTypes;
         private final List<Integer> sortChannels;
         private final List<SortOrder> sortOrders;
@@ -241,12 +240,14 @@ private static class TopNBuilder
         private long memorySize;
 
         private TopNBuilder(int n,
+                boolean partial,
                 List<Type> sortTypes,
                 List<Integer> sortChannels,
                 List<SortOrder> sortOrders,
                 OperatorContext operatorContext)
         {
             this.n = n;
+            this.partial = partial;
 
             this.sortTypes = sortTypes;
             this.sortChannels = sortChannels;
@@ -336,7 +337,16 @@ private boolean isFull()
             if (memorySize < 0) {
                 memorySize = 0;
             }
-            return operatorContext.setMemoryReservation(memorySize, true) != memorySize;
+            try {
+                operatorContext.setMemoryReservation(memorySize);
+                return false;
+            }
+            catch (ExceededMemoryLimitException e) {
+                if (partial) {
+                    return true;
+                }
+                throw e;
+            }
         }
 
         public Iterator<Block[]> build()