From 01e278afec16c80cd6f98b2794b3a72070d4ac56 Mon Sep 17 00:00:00 2001
From: adeneche <adeneche@gmail.com>
Date: Mon, 13 Jul 2015 09:09:40 -0700
Subject: [PATCH] DRILL-3616: Memory leak in a cleanup code after canceling
 queries with window functions spilling to disk

---
 .../impl/xsort/ExternalSortBatch.java         | 29 ++++++++++++++-----
 1 file changed, 22 insertions(+), 7 deletions(-)

diff --git a/exec/java-exec/src/main/java/org/apache/drill/exec/physical/impl/xsort/ExternalSortBatch.java b/exec/java-exec/src/main/java/org/apache/drill/exec/physical/impl/xsort/ExternalSortBatch.java
index 4bb15728fdc..67158e67602 100644
--- a/exec/java-exec/src/main/java/org/apache/drill/exec/physical/impl/xsort/ExternalSortBatch.java
+++ b/exec/java-exec/src/main/java/org/apache/drill/exec/physical/impl/xsort/ExternalSortBatch.java
@@ -20,6 +20,7 @@
 import io.netty.buffer.DrillBuf;
 
 import java.io.IOException;
+import java.util.Collection;
 import java.util.Iterator;
 import java.util.LinkedList;
 import java.util.List;
@@ -153,17 +154,31 @@ public SelectionVector4 getSelectionVector4() {
     return sv4;
   }
 
+  private void cleanupBatchGroups(Collection<BatchGroup> groups) {
+    for (BatchGroup group: groups) {
+      try {
+        group.cleanup();
+      } catch (IOException e) {
+        // collect all failure and make sure to cleanup all remaining batches
+        // Originally we would have thrown a RuntimeException that would propagate to FragmentExecutor.closeOutResources()
+        // where it would have been passed to context.fail()
+        // passing the exception directly to context.fail(e) will let the cleanup process continue instead of stopping
+        // right away, this will also make sure we collect any additional exception we may get while cleaning up
+        context.fail(e);
+      }
+    }
+  }
+
   @Override
   public void close() {
     try {
       if (batchGroups != null) {
-        for (BatchGroup group: batchGroups) {
-          try {
-            group.cleanup();
-          } catch (IOException e) {
-            throw new RuntimeException(e);
-          }
-        }
+        cleanupBatchGroups(batchGroups);
+        batchGroups = null;
+      }
+      if (spilledBatchGroups != null) {
+        cleanupBatchGroups(spilledBatchGroups);
+        spilledBatchGroups = null;
       }
     } finally {
       if (builder != null) {