[CARMEL-7548][CARMEL-4905][CARMEL-5014] Remove the leaked live entities with a TTL map (apache#292)

Author: wangyum

core/src/main/scala/org/apache/spark/internal/config/Status.scala

@@ -82,4 +82,25 @@ private[spark] object Status {
     .version("3.4.0")
     .stringConf
     .createOptional
+
+  val LIVE_JOBS_EXPIRE_HOURS = ConfigBuilder("spark.ui.liveJobs.expireHours")
+    .version("3.5.0")
+    .doc("The TTL of live jobs in AppStatusListener. Default value is -1, means the live " +
+      "jobs are never expired.")
+    .intConf
+    .createWithDefault(-1)
+
+  val LIVE_STAGES_EXPIRE_HOURS = ConfigBuilder("spark.ui.liveStages.expireHours")
+    .version("3.5.0")
+    .doc("The TTL of live stages in AppStatusListener. Default value equals to the value of " +
+      s"${LIVE_JOBS_EXPIRE_HOURS.key}.")
+    .intConf
+    .createWithDefault(LIVE_JOBS_EXPIRE_HOURS.defaultValue.get)
+
+  val LIVE_TASKS_EXPIRE_HOURS = ConfigBuilder("spark.ui.liveTasks.expireHours")
+    .version("3.5.0")
+    .doc("The TTL of live tasks in AppStatusListener. Default value equals to the value of " +
+      s"${LIVE_STAGES_EXPIRE_HOURS.key}.")
+    .intConf
+    .createWithDefault(LIVE_STAGES_EXPIRE_HOURS.defaultValue.get)
 }

core/src/main/scala/org/apache/spark/status/AppStatusListener.scala

@@ -18,13 +18,15 @@
 package org.apache.spark.status
 
 import java.util.Date
-import java.util.concurrent.ConcurrentHashMap
+import java.util.concurrent.{ConcurrentHashMap, TimeUnit}
 
 import scala.collection.JavaConverters._
 import scala.collection.concurrent.TrieMap
 import scala.collection.mutable
 import scala.collection.mutable.ArrayBuffer
 
+import com.google.common.cache.{CacheBuilder, RemovalListener, RemovalNotification}
+
 import org.apache.spark._
 import org.apache.spark.executor.{ExecutorMetrics, TaskMetrics}
 import org.apache.spark.internal.Logging
@@ -73,11 +75,61 @@ private[spark] class AppStatusListener(
 
   // Keep track of live entities, so that task metrics can be efficiently updated (without
   // causing too many writes to the underlying store, and other expensive operations).
-  private val liveStages = new ConcurrentHashMap[(Int, Int), LiveStage]()
-  private val liveJobs = new TrieMap[Int, LiveJob]()
+
+  // In a long running thrift-server, events may be dropped due to limited event queue.
+  // If JobEnd/StageCompleted/TaskEnd events missed, the LiveJob/LiveStage/LiveTask instances
+  // will be kept in memory forever. To prevent this leak, use a map with expired time (TTL)
+  // to remove these leak instances automatically.
+  private val liveJobs =
+  if (conf.get(LIVE_JOBS_EXPIRE_HOURS) <= 0) {
+    new ConcurrentHashMap[java.lang.Integer, LiveJob]()
+  } else {
+    CacheBuilder.newBuilder()
+      .expireAfterWrite(conf.get(LIVE_JOBS_EXPIRE_HOURS), TimeUnit.HOURS)
+      .removalListener(new RemovalListener[java.lang.Integer, LiveJob] {
+        override def onRemoval(rn: RemovalNotification[Integer, LiveJob]): Unit = {
+          // We set job SUCCEEDED, when the trigger of kvstore invokes cleanupJobs,
+          // it will release from kvstore to avoid memory leak.
+          val job = rn.getValue
+          job.status = JobExecutionStatus.SUCCEEDED
+          update(job, System.nanoTime(), last = true)
+          logInfo(s"Try to clean live job ${rn.getKey} and update its status in kvstore")
+        }
+      }).build[java.lang.Integer, LiveJob]().asMap()
+  }
+  private val liveStages =
+    if (conf.get(LIVE_STAGES_EXPIRE_HOURS) <= 0) {
+      new ConcurrentHashMap[(java.lang.Integer, java.lang.Integer), LiveStage]()
+    } else {
+      CacheBuilder.newBuilder()
+        .expireAfterWrite(conf.get(LIVE_STAGES_EXPIRE_HOURS), TimeUnit.HOURS)
+        .removalListener(new RemovalListener[(java.lang.Integer, java.lang.Integer), LiveStage] {
+          override def onRemoval(rn: RemovalNotification[(Integer, Integer), LiveStage]): Unit = {
+            val stage = rn.getValue
+            stage.cleaning = true
+            // cleanupTasks is heavy, async
+            kvstore.doAsync {
+              cleanupTasks(stage, skipActive = false)
+              // We set stage COMPLETE here, when the trigger of kvstore invokes cleanupStages,
+              // it will release from kvstore to avoid memory leak.
+              stage.status == v1.StageStatus.COMPLETE
+              update(stage, System.nanoTime(), last = true)
+              logInfo(s"Try to clean live stage ${rn.getKey._1}.${rn.getKey._2} and " +
+                s"update its status in kvstore")
+            }
+          }
+        }).build[(java.lang.Integer, java.lang.Integer), LiveStage]().asMap()
+    }
+  private val liveTasks =
+    if (conf.get(LIVE_TASKS_EXPIRE_HOURS) <= 0) {
+      new ConcurrentHashMap[java.lang.Long, LiveTask]()
+    } else {
+      CacheBuilder.newBuilder()
+        .expireAfterWrite(conf.get(LIVE_TASKS_EXPIRE_HOURS), TimeUnit.HOURS)
+        .build[java.lang.Long, LiveTask]().asMap()
+    }
   private[spark] val liveExecutors = new TrieMap[String, LiveExecutor]()
   private[spark] val deadExecutors = new TrieMap[String, LiveExecutor]()
-  private val liveTasks = new TrieMap[Long, LiveTask]()
   private val liveRDDs = new TrieMap[Int, LiveRDD]()
   private val pools = new TrieMap[String, SchedulerPool]()
   private val liveResourceProfiles = new TrieMap[Int, LiveResourceProfile]()
@@ -494,7 +546,7 @@ private[spark] class AppStatusListener(
   }
 
   override def onJobEnd(event: SparkListenerJobEnd): Unit = {
-    liveJobs.remove(event.jobId).foreach { job =>
+    Option(liveJobs.remove(event.jobId)).foreach { job =>
       val now = System.nanoTime()
 
       // Check if there are any pending stages that match this job; mark those as skipped.
@@ -565,7 +617,7 @@ private[spark] class AppStatusListener(
     }.getOrElse(SparkUI.DEFAULT_POOL_NAME)
 
     // Look at all active jobs to find the ones that mention this stage.
-    stage.jobs = liveJobs.values
+    stage.jobs = liveJobs.values.asScala
       .filter(_.stageIds.contains(event.stageInfo.stageId))
       .toSeq
     stage.jobIds = stage.jobs.map(_.jobId).toSet
@@ -638,7 +690,7 @@ private[spark] class AppStatusListener(
   override def onTaskGettingResult(event: SparkListenerTaskGettingResult): Unit = {
     // Call update on the task so that the "getting result" time is written to the store; the
     // value is part of the mutable TaskInfo state that the live entity already references.
-    liveTasks.get(event.taskInfo.taskId).foreach { task =>
+    Option(liveTasks.get(event.taskInfo.taskId)).foreach { task =>
       maybeUpdate(task, System.nanoTime())
     }
   }
@@ -651,7 +703,7 @@ private[spark] class AppStatusListener(
 
     val now = System.nanoTime()
 
-    val metricsDelta = liveTasks.remove(event.taskInfo.taskId).map { task =>
+    val metricsDelta = Option(liveTasks.remove(event.taskInfo.taskId)).map { task =>
       task.info = event.taskInfo
 
       val errorMessage = event.reason match {
@@ -941,7 +993,7 @@ private[spark] class AppStatusListener(
     val now = System.nanoTime()
 
     event.accumUpdates.foreach { case (taskId, sid, sAttempt, accumUpdates) =>
-      liveTasks.get(taskId).foreach { task =>
+      Option(liveTasks.get(taskId)).foreach { task =>
         val metrics = TaskMetrics.fromAccumulatorInfos(accumUpdates)
         val delta = task.updateMetrics(metrics)
         maybeUpdate(task, now)
@@ -1032,9 +1084,9 @@ private[spark] class AppStatusListener(
       entityFlushFunc(stage)
       stage.executorSummaries.values.foreach(entityFlushFunc)
     }
-    liveJobs.values.foreach(entityFlushFunc)
+    liveJobs.values.forEach(e => entityFlushFunc(e))
     liveExecutors.values.foreach(entityFlushFunc)
-    liveTasks.values.foreach(entityFlushFunc)
+    liveTasks.values.forEach(e => entityFlushFunc(e))
     liveRDDs.values.foreach(entityFlushFunc)
     pools.values.foreach(entityFlushFunc)
   }
@@ -1213,7 +1265,7 @@ private[spark] class AppStatusListener(
 
   private def getOrCreateStage(info: StageInfo): LiveStage = {
     val stage = liveStages.computeIfAbsent((info.stageId, info.attemptNumber),
-      (_: (Int, Int)) => new LiveStage(info))
+      (_: (java.lang.Integer, java.lang.Integer)) => new LiveStage(info))
     stage.info = info
     stage
   }
@@ -1274,7 +1326,12 @@ private[spark] class AppStatusListener(
     val toDelete = KVUtils.viewToSeq(view, countToDelete.toInt) { j =>
       j.info.status != JobExecutionStatus.RUNNING && j.info.status != JobExecutionStatus.UNKNOWN
     }
-    toDelete.foreach { j => kvstore.delete(j.getClass(), j.info.jobId) }
+    toDelete.foreach { j =>
+      if (j.info.status == JobExecutionStatus.SUCCEEDED && j.info.completionTime.isEmpty) {
+        logInfo(s"Clean up leaked job ${j.info.jobId} from kvstore")
+      }
+      kvstore.delete(j.getClass(), j.info.jobId)
+    }
   }
 
   private case class StageCompletionTime(
@@ -1331,6 +1388,9 @@ private[spark] class AppStatusListener(
       s"including ${stages.map(_.info.stageId).mkString(",")}.")
 
     stages.map { s =>
+      if (s.info.status == v1.StageStatus.COMPLETE && s.info.completionTime.isEmpty) {
+        logInfo(s"Clean up leaked stage ${s.info.stageId}.${s.info.attemptId} from kvstore")
+      }
       val key = Array(s.info.stageId, s.info.attemptId)
       kvstore.delete(s.getClass(), key)
 
@@ -1381,7 +1441,7 @@ private[spark] class AppStatusListener(
     kvstore.removeAllByIndexValues(classOf[TaskDataWrapper], TaskIndexNames.STAGE, stageIds)
   }
 
-  private def cleanupTasks(stage: LiveStage): Unit = {
+  private def cleanupTasks(stage: LiveStage, skipActive: Boolean = true): Unit = {
     val countToDelete = calculateNumberToRemove(stage.savedTasks.get(), maxTasksPerStage).toInt
     if (countToDelete > 0) {
       val stageKey = Array(stage.info.stageId, stage.info.attemptNumber)
@@ -1391,7 +1451,7 @@ private[spark] class AppStatusListener(
 
       // Try to delete finished tasks only.
       val toDelete = KVUtils.viewToSeq(view, countToDelete) { t =>
-        !live || t.status != TaskState.RUNNING.toString()
+        !live || (skipActive && t.status != TaskState.RUNNING.toString())
       }
       toDelete.foreach { t => kvstore.delete(t.getClass(), t.taskId) }
       stage.savedTasks.addAndGet(-toDelete.size)