Fix memory leak in MemoryAwareQueryExecution

Previously, we would add a listener for state change *on* every state change. This could also cause listeners to be added in a loop in between failed state transitions. Fixes prestodb#10812.
raghavsethi · Jul 23, 2018 · 4cb93d5 · 4cb93d5
1 parent 3b50e03
commit 4cb93d5
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 9 deletions.
diff --git a/presto-main/src/main/java/com/facebook/presto/execution/MemoryAwareQueryExecution.java b/presto-main/src/main/java/com/facebook/presto/execution/MemoryAwareQueryExecution.java
@@ -43,6 +43,9 @@ public class MemoryAwareQueryExecution
     @GuardedBy("this")
     private boolean startedWaiting;
 
+    @GuardedBy("this")
+    private boolean preAllocationSucceeded;
+
     public MemoryAwareQueryExecution(ClusterMemoryManager memoryManager, SqlQueryExecution delegate)
     {
         this.memoryManager = memoryManager;
@@ -139,14 +142,17 @@ public synchronized void start()
     {
         try (SetThreadName ignored = new SetThreadName("Query-%s", delegate.getQueryId())) {
             try {
-                if (memoryManager.preAllocateQueryMemory(delegate.getQueryId(), peakMemoryEstimate)) {
-                    delegate.addStateChangeListener(state -> {
-                        if (state.isDone()) {
-                            memoryManager.removePreAllocation(delegate.getQueryId());
-                        }
-                    });
-                    delegate.start();
-                    return;
+                if (!preAllocationSucceeded) {
+                    if (memoryManager.preAllocateQueryMemory(delegate.getQueryId(), peakMemoryEstimate)) {
+                        preAllocationSucceeded = true;
+                        delegate.addStateChangeListener(state -> {
+                            if (state.isDone()) {
+                                memoryManager.removePreAllocation(delegate.getQueryId());
+                            }
+                        });
+                        delegate.start();
+                        return;
+                    }
                 }
 
                 if (!startedWaiting) {

diff --git a/presto-main/src/main/java/com/facebook/presto/memory/ClusterMemoryManager.java b/presto-main/src/main/java/com/facebook/presto/memory/ClusterMemoryManager.java
@@ -107,7 +107,10 @@ public class ClusterMemoryManager
     private final AtomicLong clusterMemoryBytes = new AtomicLong();
     private final AtomicLong queriesKilledDueToOutOfMemory = new AtomicLong();
 
-    private final Map<QueryId, Long> preAllocations = new HashMap<>();
+    @GuardedBy("this")
+    public final Map<QueryId, Long> preAllocations = new HashMap<>();
+
+    @GuardedBy("this")
     private final Map<QueryId, Long> preAllocationsConsumed = new HashMap<>();
 
     @GuardedBy("this")