[SPARK-22713][CORE] ExternalAppendOnlyMap leaks when spilled during iteration

core/src/main/scala/org/apache/spark/util/collection/ExternalAppendOnlyMap.scala

@@ -80,7 +80,10 @@ class ExternalAppendOnlyMap[K, V, C](
     this(createCombiner, mergeValue, mergeCombiners, serializer, blockManager, TaskContext.get())
   }
 
-  @volatile private var currentMap = new SizeTrackingAppendOnlyMap[K, C]
+  /**
+   * Exposed for testing
+   */
+  @volatile private[collection] var currentMap = new SizeTrackingAppendOnlyMap[K, C]
   private val spilledMaps = new ArrayBuffer[DiskMapIterator]
   private val sparkConf = SparkEnv.get.conf
   private val diskBlockManager = blockManager.diskBlockManager
@@ -114,7 +117,10 @@ class ExternalAppendOnlyMap[K, V, C](
   private val keyComparator = new HashComparator[K]
   private val ser = serializer.newInstance()
 
-  @volatile private var readingIterator: SpillableIterator = null
+  /**
+   * Exposed for testing
+   */
+  @volatile private[collection] var readingIterator: SpillableIterator = null
 
   /**
    * Number of files this map has spilled so far.
@@ -267,7 +273,7 @@ class ExternalAppendOnlyMap[K, V, C](
    */
   def destructiveIterator(inMemoryIterator: Iterator[(K, C)]): Iterator[(K, C)] = {
     readingIterator = new SpillableIterator(inMemoryIterator)
-    readingIterator
+    readingIterator.toCompletionIterator
   }
 
   /**
@@ -280,8 +286,7 @@ class ExternalAppendOnlyMap[K, V, C](
         "ExternalAppendOnlyMap.iterator is destructive and should only be called once.")
     }
     if (spilledMaps.isEmpty) {
-      CompletionIterator[(K, C), Iterator[(K, C)]](
-        destructiveIterator(currentMap.iterator), freeCurrentMap())
+      destructiveIterator(currentMap.iterator)
     } else {
       new ExternalIterator()
     }
@@ -305,8 +310,8 @@ class ExternalAppendOnlyMap[K, V, C](
 
     // Input streams are derived both from the in-memory map and spilled maps on disk
     // The in-memory map is sorted in place, while the spilled maps are already in sorted order
-    private val sortedMap = CompletionIterator[(K, C), Iterator[(K, C)]](destructiveIterator(
-      currentMap.destructiveSortedIterator(keyComparator)), freeCurrentMap())
+    private val sortedMap = destructiveIterator(
+      currentMap.destructiveSortedIterator(keyComparator))
     private val inputStreams = (Seq(sortedMap) ++ spilledMaps).map(it => it.buffered)
 
     inputStreams.foreach { it =>
@@ -568,13 +573,14 @@ class ExternalAppendOnlyMap[K, V, C](
     context.addTaskCompletionListener(context => cleanup())
   }
 
-  private[this] class SpillableIterator(var upstream: Iterator[(K, C)])
+  /**
+   * Exposed for testing
+   */
+  private[collection] class SpillableIterator(var upstream: Iterator[(K, C)])
     extends Iterator[(K, C)] {
 
     private val SPILL_LOCK = new Object()
 
-    private var nextUpstream: Iterator[(K, C)] = null
-
     private var cur: (K, C) = readNext()
 
     private var hasSpilled: Boolean = false
@@ -585,17 +591,25 @@ class ExternalAppendOnlyMap[K, V, C](
       } else {
         logInfo(s"Task ${context.taskAttemptId} force spilling in-memory map to disk and " +
           s"it will release ${org.apache.spark.util.Utils.bytesToString(getUsed())} memory")
-        nextUpstream = spillMemoryIteratorToDisk(upstream)
+        val nextUpstream = spillMemoryIteratorToDisk(upstream)
+        assert(!upstream.hasNext)
         hasSpilled = true
+        upstream = nextUpstream
         true
       }
     }
 
+    private def destroy() : Unit = {
+      freeCurrentMap()
+      upstream = Iterator.empty
+    }
+
+    private[ExternalAppendOnlyMap]
+    def toCompletionIterator: CompletionIterator[(K, C), SpillableIterator] = {
+      CompletionIterator[(K, C), SpillableIterator](this, this.destroy )
+    }
+
     def readNext(): (K, C) = SPILL_LOCK.synchronized {
-      if (nextUpstream != null) {
-        upstream = nextUpstream
-        nextUpstream = null
-      }
       if (upstream.hasNext) {
         upstream.next()
       } else {

core/src/test/scala/org/apache/spark/util/collection/ExternalAppendOnlyMapSuite.scala

@@ -23,6 +23,7 @@ import org.apache.spark._
 import org.apache.spark.internal.config._
 import org.apache.spark.io.CompressionCodec
 import org.apache.spark.memory.MemoryTestingUtils
+import org.apache.spark.util.CompletionIterator
 
 class ExternalAppendOnlyMapSuite extends SparkFunSuite with LocalSparkContext {
   import TestUtils.{assertNotSpilled, assertSpilled}
@@ -414,6 +415,99 @@ class ExternalAppendOnlyMapSuite extends SparkFunSuite with LocalSparkContext {
     sc.stop()
   }
 
+  test("spill during iteration") {
+    val size = 1000
+    val conf = createSparkConf(loadDefaults = true)
+    sc = new SparkContext("local-cluster[1,1,1024]", "test", conf)
+    val map = createExternalMap[Int]
+
+    map.insertAll((0 until size).iterator.map(i => (i / 10, i)))
+    assert(map.numSpills == 0, "map was not supposed to spill")
+
+    val it = map.iterator
+    assert( it.isInstanceOf[CompletionIterator[_, _]])
+    val underlyingIt = map.readingIterator
+    assert( underlyingIt != null )
+    val underlyingMapIterator = underlyingIt.upstream
+    assert(underlyingMapIterator != null)
+    val underlyingMapIteratorClass = underlyingMapIterator.getClass
+    assert(underlyingMapIteratorClass.getEnclosingClass == classOf[AppendOnlyMap[_, _]])
+
+    val underlyingMap = map.currentMap
+    assert(underlyingMap != null)
+
+    val first50Keys = for ( _ <- 0 until 50) yield {
+      val (k, vs) = it.next
+      val sortedVs = vs.sorted
+      assert(sortedVs.seq == (0 until 10).map(10 * k + _))
+      k
+    }
+    assert( map.numSpills == 0 )
+    map.spill(Long.MaxValue, null)
+    // these asserts try to show that we're no longer holding references to the underlying map.
+    // it'd be nice to use something like
+    // https://github.com/scala/scala/blob/2.13.x/test/junit/scala/tools/testing/AssertUtil.scala
+    // (lines 69-89)
+    assert(map.currentMap == null)
+    assert(underlyingIt.upstream ne underlyingMapIterator)
+    assert(underlyingIt.upstream.getClass != underlyingMapIteratorClass)
+    assert(underlyingIt.upstream.getClass.getEnclosingClass != classOf[AppendOnlyMap[_, _]])
+
+    val next50Keys = for ( _ <- 0 until 50) yield {
+      val (k, vs) = it.next
+      val sortedVs = vs.sorted
+      assert(sortedVs.seq == (0 until 10).map(10 * k + _))
+      k
+    }
+    assert(!it.hasNext)
+    val keys = (first50Keys ++ next50Keys).sorted
+    assert(keys == (0 until 100))
+  }
+
+  test("drop all references to the underlying map once the iterator is exhausted") {
+    val size = 1000
+    val conf = createSparkConf(loadDefaults = true)
+    sc = new SparkContext("local-cluster[1,1,1024]", "test", conf)
+    val map = createExternalMap[Int]
+
+    map.insertAll((0 until size).iterator.map(i => (i / 10, i)))
+    assert(map.numSpills == 0, "map was not supposed to spill")
+
+    val it = map.iterator
+    assert( it.isInstanceOf[CompletionIterator[_, _]])
+    val underlyingIt = map.readingIterator
+    assert( underlyingIt != null )
+    val underlyingMapIterator = underlyingIt.upstream
+    assert(underlyingMapIterator != null)
+    val underlyingMapIteratorClass = underlyingMapIterator.getClass
+    assert(underlyingMapIteratorClass.getEnclosingClass == classOf[AppendOnlyMap[_, _]])
+
+    val underlyingMap = map.currentMap
+    assert(underlyingMap != null)
+
+    val keys = it.map{
+      case (k, vs) =>
+        val sortedVs = vs.sorted
+        assert(sortedVs.seq == (0 until 10).map(10 * k + _))
+        k
+    }
+    .toList
+    .sorted
+
+    assert(it.isEmpty)
+    assert(keys == (0 until 100))
+
+    assert( map.numSpills == 0 )
+    // these asserts try to show that we're no longer holding references to the underlying map.
+    // it'd be nice to use something like
+    // https://github.com/scala/scala/blob/2.13.x/test/junit/scala/tools/testing/AssertUtil.scala
+    // (lines 69-89)
+    assert(map.currentMap == null)
+    assert(underlyingIt.upstream ne underlyingMapIterator)
+    assert(underlyingIt.upstream.getClass != underlyingMapIteratorClass)
+    assert(underlyingIt.upstream.getClass.getEnclosingClass != classOf[AppendOnlyMap[_, _]])
+  }
+
   test("external aggregation updates peak execution memory") {
     val spillThreshold = 1000
     val conf = createSparkConf(loadDefaults = false)