[SPARK-14363] Fix executor OOM due to memory leak in the Sorter

Sital Kedia · Sital Kedia · commit d89adf8e514a · 2016-04-12T13:53:55.000-07:00
diff --git a/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleExternalSorter.java b/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleExternalSorter.java
@@ -215,8 +215,6 @@ private void writeSortedFile(boolean isLastFile) throws IOException {
       }
     }
 
-    inMemSorter.reset();
-
     if (!isLastFile) {  // i.e. this is a spill file
       // The current semantics of `shuffleRecordsWritten` seem to be that it's updated when records
       // are written to disk, not when they enter the shuffle sorting code. DiskBlockObjectWriter
@@ -255,6 +253,10 @@ public long spill(long size, MemoryConsumer trigger) throws IOException {
 
     writeSortedFile(false);
     final long spillSize = freeMemory();
+    inMemSorter.reset();
+    // Reset the in-memory sorter's pointer array only after freeing up the memory pages holding the
+    // records. Otherwise, if the task is over allocated memory, then without freeing the memory pages,
+    // we might not be able to get memory for the pointer array.
     taskContext.taskMetrics().incMemoryBytesSpilled(spillSize);
     return spillSize;
   }
diff --git a/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleInMemorySorter.java b/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleInMemorySorter.java
@@ -51,9 +51,12 @@ public int compare(PackedRecordPointer left, PackedRecordPointer right) {
    */
   private int pos = 0;
 
+  private int initialSize;
+
   ShuffleInMemorySorter(MemoryConsumer consumer, int initialSize) {
     this.consumer = consumer;
     assert (initialSize > 0);
+    this.initialSize = initialSize;
     this.array = consumer.allocateArray(initialSize);
     this.sorter = new Sorter<>(ShuffleSortDataFormat.INSTANCE);
   }
@@ -70,6 +73,10 @@ public int numRecords() {
   }
 
   public void reset() {
+    if (consumer != null) {
+      consumer.freeArray(array);
+      this.array = consumer.allocateArray(initialSize);
+    }
     pos = 0;
   }
 
diff --git a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
@@ -200,14 +200,17 @@ public long spill(long size, MemoryConsumer trigger) throws IOException {
         spillWriter.write(baseObject, baseOffset, recordLength, sortedRecords.getKeyPrefix());
       }
       spillWriter.close();
-
-      inMemSorter.reset();
     }
 
     final long spillSize = freeMemory();
     // Note that this is more-or-less going to be a multiple of the page size, so wasted space in
     // pages will currently be counted as memory spilled even though that space isn't actually
     // written to disk. This also counts the space needed to store the sorter's pointer array.
+    inMemSorter.reset();
+    // Reset the in-memory sorter's pointer array only after freeing up the memory pages holding the
+    // records. Otherwise, if the task is over allocated memory, then without freeing the memory pages,
+    // we might not be able to get memory for the pointer array.
+
     taskContext.taskMetrics().incMemoryBytesSpilled(spillSize);
 
     return spillSize;
diff --git a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorter.java b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorter.java
@@ -84,6 +84,8 @@ public int compare(RecordPointerAndKeyPrefix r1, RecordPointerAndKeyPrefix r2) {
    */
   private int pos = 0;
 
+  private long initialSize;
+
   public UnsafeInMemorySorter(
     final MemoryConsumer consumer,
     final TaskMemoryManager memoryManager,
@@ -102,6 +104,7 @@ public UnsafeInMemorySorter(
       LongArray array) {
     this.consumer = consumer;
     this.memoryManager = memoryManager;
+    this.initialSize = array.size();
     if (recordComparator != null) {
       this.sorter = new Sorter<>(UnsafeSortDataFormat.INSTANCE);
       this.sortComparator = new SortComparator(recordComparator, prefixComparator, memoryManager);
@@ -123,6 +126,10 @@ public void free() {
   }
 
   public void reset() {
+    if (consumer != null) {
+      consumer.freeArray(array);
+      this.array = consumer.allocateArray(initialSize);
+    }
     pos = 0;
   }
 

Original file line number	Diff line number	Diff line change
`@@ -215,8 +215,6 @@ private void writeSortedFile(boolean isLastFile) throws IOException {`
`215`	`215`	`}`
`216`	`216`	`}`
`217`	`217`
`218`		`- inMemSorter.reset();`
`219`		`-`
`220`	`218`	`if (!isLastFile) { // i.e. this is a spill file`
`221`	`219`	// The current semantics of `shuffleRecordsWritten` seem to be that it's updated when records
`222`	`220`	`// are written to disk, not when they enter the shuffle sorting code. DiskBlockObjectWriter`
`@@ -255,6 +253,10 @@ public long spill(long size, MemoryConsumer trigger) throws IOException {`
`255`	`253`
`256`	`254`	`writeSortedFile(false);`
`257`	`255`	`final long spillSize = freeMemory();`
	`256`	`+ inMemSorter.reset();`
	`257`	`+ // Reset the in-memory sorter's pointer array only after freeing up the memory pages holding the`
	`258`	`+ // records. Otherwise, if the task is over allocated memory, then without freeing the memory pages,`
	`259`	`+ // we might not be able to get memory for the pointer array.`
`258`	`260`	`taskContext.taskMetrics().incMemoryBytesSpilled(spillSize);`
`259`	`261`	`return spillSize;`
`260`	`262`	`}`