[SPARK-17503][Core] Fix memory leak in Memory store when unable to cache the whole RDD in memory

[clockfly]

What changes were proposed in this pull request?

MemoryStore may throw OutOfMemoryError when trying to cache a super big RDD that cannot fit in memory.

   scala> sc.parallelize(1 to 1000000000, 100).map(x => new Array[Long](1000)).cache().count()

   java.lang.OutOfMemoryError: Java heap space
    at $line14.$read$$iw$$iw$$iw$$iw$$iw$$iw$$iw$$iw$$anonfun$1.apply(<console>:24)
    at $line14.$read$$iw$$iw$$iw$$iw$$iw$$iw$$iw$$iw$$anonfun$1.apply(<console>:23)
    at scala.collection.Iterator$$anon$11.next(Iterator.scala:409)
    at scala.collection.Iterator$JoinIterator.next(Iterator.scala:232)
    at org.apache.spark.storage.memory.PartiallyUnrolledIterator.next(MemoryStore.scala:683)
    at org.apache.spark.InterruptibleIterator.next(InterruptibleIterator.scala:43)
    at org.apache.spark.util.Utils$.getIteratorSize(Utils.scala:1684)
    at org.apache.spark.rdd.RDD$$anonfun$count$1.apply(RDD.scala:1134)
    at org.apache.spark.rdd.RDD$$anonfun$count$1.apply(RDD.scala:1134)
    at org.apache.spark.SparkContext$$anonfun$runJob$5.apply(SparkContext.scala:1915)
    at org.apache.spark.SparkContext$$anonfun$runJob$5.apply(SparkContext.scala:1915)
    at org.apache.spark.scheduler.ResultTask.runTask(ResultTask.scala:70)
    at org.apache.spark.scheduler.Task.run(Task.scala:86)
    at org.apache.spark.executor.Executor$TaskRunner.run(Executor.scala:274)
    at java.util.concurrent.ThreadPoolExecutor.runWorker(ThreadPoolExecutor.java:1142)
    at java.util.concurrent.ThreadPoolExecutor$Worker.run(ThreadPoolExecutor.java:617)
    at java.lang.Thread.run(Thread.java:745)

Spark MemoryStore uses SizeTrackingVector as a temporary unrolling buffer to store all input values that it has read so far before transferring the values to storage memory cache. The problem is that when the input RDD is too big for caching in memory, the temporary unrolling memory SizeTrackingVector is not garbage collected in time. As SizeTrackingVector can occupy all available storage memory, it may cause the executor JVM to run out of memory quickly.

More info can be found at https://issues.apache.org/jira/browse/SPARK-17503

How was this patch tested?

Unit test and manual test.

Before change

Heap memory consumption

Heap dump

After change

Heap memory consumption

[clockfly]

@yhuai

[SparkQA]

Test build #65243 has finished for PR 15056 at commit a9a4a8b.

• This patch passes all tests.
• This patch merges cleanly.
• This patch adds no public classes.

[srowen]

Oh I get it now. That looks good, to my understanding.

[cloud-fan]

LGTM

[yhuai]

Let's see if I understand the problem.

Because BlockManager may call close early, so we cannot rely on CompletionIterator to free the memory because we will never actually consume all of elements of unrolled, right?

[JoshRosen]

I think the problem here is that the completion iterator is releasing the bookkeeping memory for the iterator as soon as the iterator is fully iterated, but the on-heap objects are being retained by the reference in the unrolled field.

[yhuai]

oh, yea, you are right.

[JoshRosen]

LGTM as well, so I'm going to merge this to master and branch-2.0 (2.0.1-SNAPSHOT). Thanks!