[CARMEL-7382][CARMEL-6292] Fix spilled query result leak (apache#122)

Author: wangyum

core/src/main/scala/org/apache/spark/ContextCleaner.scala

@@ -485,7 +485,7 @@ private[spark] class CleanSparkListenerCleanupWorker(sc: SparkContext)
 private[spark] class SpilledPartitionResultCleanupWorker(sc: SparkContext)
   extends ContextCleanupWorker(sc, classOf[CleanSpilledPartitionResult].getName) with Logging {
 
-  private val deleteCount = new AtomicLong(0L)
+  private[spark] val deleteCount = new AtomicLong(0L)
 
   override def doCleanup(task: CleanupTask): Unit = {
     task match {

core/src/main/scala/org/apache/spark/scheduler/AnalyticsTaskSchedulerImpl.scala

@@ -290,6 +290,7 @@ private[spark] class AnalyticsTaskSchedulerImpl(
           taskSetsByStageIdAndAttempt -= manager.taskSet.stageId
         }
       }
+      taskResultGetter.cleanTaskSet(manager.taskSet.id)
       manager.parent.removeSchedulable(manager)
     } finally {
       manager.taskSetLock.unlock()

sql/core/src/test/scala/org/apache/spark/sql/SpillDirectResultSuite.scala

@@ -19,7 +19,7 @@ package org.apache.spark.sql
 
 import scala.collection.mutable.ArrayBuffer
 
-import org.apache.spark.{DebugFilesystem, SparkConf}
+import org.apache.spark.{CleanSpilledPartitionResult, DebugFilesystem, SparkConf, SpilledPartitionResultCleanupWorker}
 import org.apache.spark.internal.config._
 import org.apache.spark.scheduler.{SimpleRepeatableIterator, SpilledResultIterator}
 import org.apache.spark.sql.internal.SQLConf
@@ -179,4 +179,91 @@ class SpillDirectResultSuite extends QueryTest with SQLTestUtils with SharedSpar
       }
     }
   }
+
+  test("DirectResult spilled result could be cleaned by context cleaner") {
+    val spilledPartitionResultCleanupWorker = sparkContext.cleaner.get.
+      getContextCleanupWorker(classOf[CleanSpilledPartitionResult].getName).
+      asInstanceOf[SpilledPartitionResultCleanupWorker]
+
+    // val cleanCountBefore = spilledPartitionResultCleanupWorker.cleanCount.get
+    val deleteCountBefore = spilledPartitionResultCleanupWorker.deleteCount.get
+    var referenceBufferSizeBefore = 0
+
+    withTable(table2) {
+      newTable(table2, 300)
+
+      val query = s"SELECT * FROM $table2 t2 order by key"
+      val df = sql(query)
+      val it = df.collectAsIterator()
+      assert(it.isInstanceOf[SpilledResultIterator[_, _]])
+      assert(it.asInstanceOf[SpilledResultIterator[Any, Any]].rowCount == 300)
+
+      val rs = ArrayBuffer[Row]()
+      while (it.hasNext) {
+        rs.append(it.next())
+      }
+      assert(rs.length == 300)
+      var i = 0
+      rs.foreach(row => {
+        assert(row.getInt(0) == i)
+        i += 1
+      })
+
+      // Skip close iterator
+      // it.close()
+
+      referenceBufferSizeBefore = spilledPartitionResultCleanupWorker.referenceBufferSize()
+      assert(referenceBufferSizeBefore > 0)
+
+      // Clean manually when ite is still available
+      System.gc()
+
+      var tryCount = 1
+      while (spilledPartitionResultCleanupWorker.referenceBufferSize() > 0 && tryCount < 100) {
+        spilledPartitionResultCleanupWorker.clean()
+        tryCount += 1
+      }
+
+      // Check cannot clean
+      assert(spilledPartitionResultCleanupWorker.referenceBufferSize() > 0)
+    }
+
+    // Clean manually when ite is not available
+    System.gc()
+
+    var tryCount = 1
+    while (spilledPartitionResultCleanupWorker.referenceBufferSize() > 0 && tryCount < 100) {
+      spilledPartitionResultCleanupWorker.clean()
+      tryCount += 1
+    }
+
+    // val cleanCountAfter = spilledPartitionResultCleanupWorker.cleanCount.get
+    assert(spilledPartitionResultCleanupWorker.referenceBufferSize() == 0)
+    // assert( (cleanCountAfter - cleanCountBefore) == referenceBufferSizeBefore)
+    assert( (spilledPartitionResultCleanupWorker.deleteCount.get - deleteCountBefore) > 0)
+  }
+
+  test("spilled result can be cleaned properly when task fail") {
+    val spilledPartitionResultCleanupWorker = sparkContext.cleaner.get.
+      getContextCleanupWorker(classOf[CleanSpilledPartitionResult].getName).
+      asInstanceOf[SpilledPartitionResultCleanupWorker]
+    withTable(table2) {
+      newTable(table2, 30000)
+
+      val query = s"SELECT * FROM $table2 t2 order by key"
+      val df = sql(query)
+      intercept[Exception] {
+        df.collectAsIterator()
+      }
+      System.gc()
+
+      var tryCount = 1
+      while (spilledPartitionResultCleanupWorker.referenceBufferSize() > 0 && tryCount < 100) {
+        spilledPartitionResultCleanupWorker.clean()
+        tryCount += 1
+        System.gc()
+      }
+      assert(spilledPartitionResultCleanupWorker.referenceBufferSize() == 0)
+    }
+  }
 }