improve OOM handling

Author: Davies Liu

core/src/main/java/org/apache/spark/shuffle/sort/ShuffleExternalSorter.java

@@ -320,22 +320,26 @@ public void cleanupResources() {
   private void growPointerArrayIfNecessary() throws IOException {
     assert(inMemSorter != null);
     if (!inMemSorter.hasSpaceForAnotherRecord()) {
-      logger.debug("Attempting to expand sort pointer array");
       long used = inMemSorter.getMemoryUsage();
-      long needed = inMemSorter.getMemoryToExpand();
+      long needed = used + inMemSorter.getMemoryToExpand();
       try {
-        acquireMemory(used + needed);  // could trigger spilling
-        if (inMemSorter.hasSpaceForAnotherRecord()) {
-          releaseMemory(used + needed);
-        } else {
-          logger.debug("Expand sort pointer array");
+        acquireMemory(needed);  // could trigger spilling
+      } catch (OutOfMemoryError e) {
+        // should have trigger spilling
+        assert(inMemSorter.hasSpaceForAnotherRecord());
+        return;
+      }
+      // check if spilling is triggered or not
+      if (inMemSorter.hasSpaceForAnotherRecord()) {
+        releaseMemory(needed);
+      } else {
+        try {
           inMemSorter.expandPointerArray();
           releaseMemory(used);
-        }
-      } catch (OutOfMemoryError oom) {
-        // spilling should be triggered
-        if (!inMemSorter.hasSpaceForAnotherRecord()) {
-          spill();  // Just in case that JVM had run out of memory
+        } catch (OutOfMemoryError oom) {
+          // Just in case that JVM had run out of memory
+          releaseMemory(needed);
+          spill();
         }
       }
     }

core/src/main/java/org/apache/spark/shuffle/sort/ShuffleInMemorySorter.java

@@ -60,8 +60,7 @@ public void reset() {
 
   private int newLength() {
     // Guard against overflow:
-    return array.length <= Integer.MAX_VALUE / 2 ?
-      (array.length * 2) : Integer.MAX_VALUE;
+    return array.length <= Integer.MAX_VALUE / 2 ? (array.length * 2) : Integer.MAX_VALUE;
   }
 
   /**

core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java

@@ -710,13 +710,13 @@ public boolean putNewKey(Object keyBase, long keyOffset, int keyLength,
   private boolean acquireNewPage(long required) {
     try {
       currentPage = allocatePage(required);
-      dataPages.add(currentPage);
-      Platform.putInt(currentPage.getBaseObject(), currentPage.getBaseOffset(), 0);
-      pageCursor = 4;
-      return true;
     } catch (OutOfMemoryError e) {
       return false;
     }
+    dataPages.add(currentPage);
+    Platform.putInt(currentPage.getBaseObject(), currentPage.getBaseOffset(), 0);
+    pageCursor = 4;
+    return true;
   }
 
   @Override
@@ -862,7 +862,13 @@ void growAndRehash() {
     final int oldCapacity = (int) oldBitSet.capacity();
 
     // Allocate the new data structures
-    allocate(Math.min(growthStrategy.nextCapacity(oldCapacity), MAX_CAPACITY));
+    try {
+      allocate(Math.min(growthStrategy.nextCapacity(oldCapacity), MAX_CAPACITY));
+    } catch (OutOfMemoryError oom) {
+      longArray = oldLongArray;
+      bitset = oldBitSet;
+      throw oom;
+    }
 
     // Re-mask (we don't recompute the hashcode because we stored all 32 bits of it)
     for (int pos = oldBitSet.nextSetBit(0); pos >= 0; pos = oldBitSet.nextSetBit(pos + 1)) {

core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java

@@ -290,19 +290,25 @@ private void growPointerArrayIfNecessary() throws IOException {
     assert(inMemSorter != null);
     if (!inMemSorter.hasSpaceForAnotherRecord()) {
       long used = inMemSorter.getMemoryUsage();
-      long needed = inMemSorter.getMemoryToExpand();
+      long needed = used + inMemSorter.getMemoryToExpand();
       try {
-        acquireMemory(used + needed);  // could trigger spilling
-        if (inMemSorter.hasSpaceForAnotherRecord()) {
-          releaseMemory(used + needed);
-        } else {
+        acquireMemory(needed);  // could trigger spilling
+      } catch (OutOfMemoryError e) {
+        // should have trigger spilling
+        assert(inMemSorter.hasSpaceForAnotherRecord());
+        return;
+      }
+      // check if spilling is triggered or not
+      if (inMemSorter.hasSpaceForAnotherRecord()) {
+        releaseMemory(needed);
+      } else {
+        try {
           inMemSorter.expandPointerArray();
           releaseMemory(used);
-        }
-      } catch (OutOfMemoryError oom) {
-        // spilling should be triggered
-        if (!inMemSorter.hasSpaceForAnotherRecord()) {
-          spill();  // Just in case that JVM had run out of memory
+        } catch (OutOfMemoryError oom) {
+          // Just in case that JVM had run out of memory
+          releaseMemory(needed);
+          spill();
         }
       }
     }