[CORE][SPARK-49788] Add optional TTL based cleaning for blocks to better support notebooks

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

@@ -19,6 +19,7 @@ package org.apache.spark
 
 import java.io.{ByteArrayInputStream, InputStream, IOException, ObjectInputStream, ObjectOutputStream}
 import java.nio.ByteBuffer
+import java.util.{HashMap => JHashMap}
 import java.util.concurrent.{ConcurrentHashMap, LinkedBlockingQueue, ThreadPoolExecutor, TimeUnit}
 import java.util.concurrent.locks.ReentrantReadWriteLock
 
@@ -711,6 +712,10 @@ private[spark] class MapOutputTrackerMaster(
     private[spark] val isLocal: Boolean)
   extends MapOutputTracker(conf) {
 
+  // Keep track of last access times for shuffle based TTL. Note: we don't use concurrent
+  // here because we don't care about overwriting times that are "close."
+  private[spark] val shuffleAccessTime = new JHashMap[Int, Long]
+
   // The size at which we use Broadcast to send the map output statuses to the executors
   private val minSizeForBroadcast = conf.get(SHUFFLE_MAPOUTPUT_MIN_SIZE_FOR_BROADCAST).toInt
 
@@ -745,6 +750,16 @@ private[spark] class MapOutputTrackerMaster(
 
   private val pushBasedShuffleEnabled = Utils.isPushBasedShuffleEnabled(conf, isDriver = true)
 
+  private[spark] val cleanerThreadpool: Option[ThreadPoolExecutor] = {
+    if (conf.get(SPARK_TTL_SHUFFLE_BLOCK_CLEANER).isDefined) {
+      val pool = ThreadUtils.newDaemonFixedThreadPool(1, "map-output-ttl-cleaner")
+      pool.execute(new TTLCleaner)
+      Some(pool)
+    } else {
+      None
+    }
+  }
+
   // Thread pool used for handling map output status requests. This is a separate thread pool
   // to ensure we don't block the normal dispatcher threads.
   private val threadpool: ThreadPoolExecutor = {
@@ -758,6 +773,68 @@ private[spark] class MapOutputTrackerMaster(
 
   private val availableProcessors = Runtime.getRuntime.availableProcessors()
 
+  def updateShuffleAtime(shuffleId: Int): Unit = {
+    if (conf.get(SPARK_TTL_SHUFFLE_BLOCK_CLEANER).isDefined) {
+      shuffleAccessTime.put(shuffleId, System.currentTimeMillis())
+    }
+  }
+
+  private class TTLCleaner extends Runnable {
+    override def run(): Unit = {
+      try {
+        // Poll the shuffle access times if we're configured for it.
+        conf.get(SPARK_TTL_SHUFFLE_BLOCK_CLEANER) match {
+          case Some(ttl) =>
+            while (true) {
+              val maxAge = System.currentTimeMillis() - ttl
+              // Find the elements to be removed & update oldest remaining time (if any)
+              var oldest = System.currentTimeMillis()
+              // Make a copy here to reduce the chance of CME
+              try {
+                val toBeRemoved = shuffleAccessTime.asScala.toList.flatMap {
+                  case (shuffleId, atime) =>
+                    if (atime < maxAge) {
+                      Some(shuffleId)
+                    } else {
+                      if (atime < oldest) {
+                        oldest = atime
+                      }
+                      None
+                    }
+                }.toList
+                toBeRemoved.map { shuffleId =>
+                  try {
+                    // Remove the shuffle access time regardless of
+                    // if we cleanup the shuffle successfully or not
+                    // since we could have a shuffle that's already
+                    // been cleaned up elsewhere.
+                    shuffleAccessTime.remove(shuffleId)
+                    unregisterAllMapAndMergeOutput(shuffleId)
+                  } catch {
+                    case NonFatal(e) =>
+                      logDebug(
+                        log"Error removing shuffle ${MDC(SHUFFLE_ID, shuffleId)}", e)
+                  }
+                }
+                // Wait until the next possible element to be removed
+                val delay = math.max((oldest + ttl) - System.currentTimeMillis(), 100)
+                Thread.sleep(delay)
+              } catch {
+                case _: java.util.ConcurrentModificationException =>
+                  // Just retry, blocks were stored while we were iterating
+                  Thread.sleep(100)
+              }
+            }
+          case None =>
+            logDebug("Tried to start TTL cleaner when not configured.")
+        }
+      } catch {
+        case _: InterruptedException =>
+          logInfo("MapOutputTrackerMaster TTLCleaner thread interrupted, exiting.")
+      }
+    }
+  }
+
   // Make sure that we aren't going to exceed the max RPC message size by making sure
   // we use broadcast to send large map output statuses.
   if (minSizeForBroadcast > maxRpcMessageSize) {
@@ -783,6 +860,7 @@ private[spark] class MapOutputTrackerMaster(
       val shuffleStatus = shuffleStatuses.get(shuffleId).head
       logDebug(s"Handling request to send ${if (needMergeOutput) "map/merge" else "map"}" +
         s" output locations for shuffle $shuffleId to $hostPort")
+      updateShuffleAtime(shuffleId)
       if (needMergeOutput) {
         context.reply(
           shuffleStatus.
@@ -834,6 +912,7 @@ private[spark] class MapOutputTrackerMaster(
   }
 
   def registerShuffle(shuffleId: Int, numMaps: Int, numReduces: Int): Unit = {
+    updateShuffleAtime(shuffleId)
     if (pushBasedShuffleEnabled) {
       if (shuffleStatuses.put(shuffleId, new ShuffleStatus(numMaps, numReduces)).isDefined) {
         throw new IllegalArgumentException("Shuffle ID " + shuffleId + " registered twice")
@@ -880,6 +959,7 @@ private[spark] class MapOutputTrackerMaster(
     shuffleStatus.removeOutputsByFilter(x => true)
     shuffleStatus.removeMergeResultsByFilter(x => true)
     shuffleStatus.removeShuffleMergerLocations()
+    shuffleAccessTime.remove(shuffleId)
     incrementEpoch()
   }
 
@@ -1257,12 +1337,14 @@ private[spark] class MapOutputTrackerMaster(
 
   // This method is only called in local-mode.
   override def getShufflePushMergerLocations(shuffleId: Int): Seq[BlockManagerId] = {
+    updateShuffleAtime(shuffleId)
     shuffleStatuses.get(shuffleId).map(_.getShufflePushMergerLocations).getOrElse(Seq.empty)
   }
 
   override def stop(): Unit = {
     mapOutputTrackerMasterMessages.offer(PoisonPill)
     threadpool.shutdown()
+    cleanerThreadpool.map(_.shutdownNow())
     try {
       sendTracker(StopMapOutputTracker)
     } catch {

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

@@ -2919,4 +2919,24 @@ package object config {
       .checkValue(v => v.forall(Set("stdout", "stderr").contains),
         "The value only can be one or more of 'stdout, stderr'.")
       .createWithDefault(Seq("stdout", "stderr"))
+
+  private[spark] val SPARK_TTL_BLOCK_CLEANER =
+    ConfigBuilder("spark.cleaner.ttl.all")
+      .doc("Add a TTL for all blocks tracked in Spark. By default blocks are only removed after " +
+        " GC on driver which with DataFrames or RDDs at the global scope will not occur. " +
+        "This must be configured before starting the SparkContext (e.g. can not be added to " +
+        "a running Spark instance.)")
+      .version("4.1.0")
+      .timeConf(TimeUnit.MILLISECONDS)
+      .createOptional
+
+  private[spark] val SPARK_TTL_SHUFFLE_BLOCK_CLEANER =
+    ConfigBuilder("spark.cleaner.ttl.shuffle")
+      .doc("Add a TTL for shuffle blocks tracked in Spark. By default blocks are only removed " +
+        "after GC on driver which with DataFrames or RDDs at the global scope will not occur." +
+        "This must be configured before starting the SparkContext (e.g. can not be added to " +
+        "a running Spark instance.)")
+      .version("4.1.0")
+      .timeConf(TimeUnit.MILLISECONDS)
+      .createOptional
 }

core/src/main/scala/org/apache/spark/storage/BlockId.scala

@@ -45,12 +45,18 @@ sealed abstract class BlockId {
     (isInstanceOf[ShuffleBlockId] || isInstanceOf[ShuffleBlockBatchId] ||
      isInstanceOf[ShuffleDataBlockId] || isInstanceOf[ShuffleIndexBlockId])
   }
+  def asShuffleId: Option[ShuffleId] = if (isShuffle) Some(asInstanceOf[ShuffleId]) else None
   def isShuffleChunk: Boolean = isInstanceOf[ShuffleBlockChunkId]
   def isBroadcast: Boolean = isInstanceOf[BroadcastBlockId]
 
   override def toString: String = name
 }
 
+@DeveloperApi
+trait ShuffleId {
+  def shuffleId: Int
+}
+
 @DeveloperApi
 case class RDDBlockId(rddId: Int, splitIndex: Int) extends BlockId {
   override def name: String = "rdd_" + rddId + "_" + splitIndex
@@ -59,7 +65,8 @@ case class RDDBlockId(rddId: Int, splitIndex: Int) extends BlockId {
 // Format of the shuffle block ids (including data and index) should be kept in sync with
 // org.apache.spark.network.shuffle.ExternalShuffleBlockResolver#getBlockData().
 @DeveloperApi
-case class ShuffleBlockId(shuffleId: Int, mapId: Long, reduceId: Int) extends BlockId {
+case class ShuffleBlockId(shuffleId: Int, mapId: Long, reduceId: Int) extends BlockId
+    with ShuffleId {
   override def name: String = "shuffle_" + shuffleId + "_" + mapId + "_" + reduceId
 }
 
@@ -69,7 +76,7 @@ case class ShuffleBlockBatchId(
     shuffleId: Int,
     mapId: Long,
     startReduceId: Int,
-    endReduceId: Int) extends BlockId {
+    endReduceId: Int) extends BlockId with ShuffleId {
   override def name: String = {
     "shuffle_" + shuffleId + "_" + mapId + "_" + startReduceId + "_" + endReduceId
   }
@@ -81,18 +88,20 @@ case class ShuffleBlockChunkId(
     shuffleId: Int,
     shuffleMergeId: Int,
     reduceId: Int,
-    chunkId: Int) extends BlockId {
+    chunkId: Int) extends BlockId with ShuffleId {
   override def name: String =
     "shuffleChunk_" + shuffleId  + "_" + shuffleMergeId + "_" + reduceId + "_" + chunkId
 }
 
 @DeveloperApi
-case class ShuffleDataBlockId(shuffleId: Int, mapId: Long, reduceId: Int) extends BlockId {
+case class ShuffleDataBlockId(shuffleId: Int, mapId: Long, reduceId: Int) extends BlockId
+    with ShuffleId {
   override def name: String = "shuffle_" + shuffleId + "_" + mapId + "_" + reduceId + ".data"
 }
 
 @DeveloperApi
-case class ShuffleIndexBlockId(shuffleId: Int, mapId: Long, reduceId: Int) extends BlockId {
+case class ShuffleIndexBlockId(shuffleId: Int, mapId: Long, reduceId: Int) extends BlockId
+    with ShuffleId {
   override def name: String = "shuffle_" + shuffleId + "_" + mapId + "_" + reduceId + ".index"
 }
 
@@ -108,7 +117,7 @@ case class ShufflePushBlockId(
     shuffleId: Int,
     shuffleMergeId: Int,
     mapIndex: Int,
-    reduceId: Int) extends BlockId {
+    reduceId: Int) extends BlockId with ShuffleId {
   override def name: String = "shufflePush_" + shuffleId + "_" +
     shuffleMergeId + "_" + mapIndex + "_" + reduceId + ""
 }
@@ -118,7 +127,7 @@ case class ShufflePushBlockId(
 case class ShuffleMergedBlockId(
     shuffleId: Int,
     shuffleMergeId: Int,
-    reduceId: Int) extends BlockId {
+    reduceId: Int) extends BlockId with ShuffleId {
   override def name: String = "shuffleMerged_" + shuffleId + "_" +
     shuffleMergeId + "_" + reduceId
 }
@@ -129,7 +138,7 @@ case class ShuffleMergedDataBlockId(
     appId: String,
     shuffleId: Int,
     shuffleMergeId: Int,
-    reduceId: Int) extends BlockId {
+    reduceId: Int) extends BlockId with ShuffleId {
   override def name: String = RemoteBlockPushResolver.MERGED_SHUFFLE_FILE_NAME_PREFIX + "_" +
     appId + "_" + shuffleId + "_" + shuffleMergeId + "_" + reduceId + ".data"
 }
@@ -140,7 +149,7 @@ case class ShuffleMergedIndexBlockId(
     appId: String,
     shuffleId: Int,
     shuffleMergeId: Int,
-    reduceId: Int) extends BlockId {
+    reduceId: Int) extends BlockId with ShuffleId {
   override def name: String = RemoteBlockPushResolver.MERGED_SHUFFLE_FILE_NAME_PREFIX + "_" +
     appId + "_" + shuffleId + "_" + shuffleMergeId + "_" + reduceId + ".index"
 }
@@ -151,7 +160,7 @@ case class ShuffleMergedMetaBlockId(
     appId: String,
     shuffleId: Int,
     shuffleMergeId: Int,
-    reduceId: Int) extends BlockId {
+    reduceId: Int) extends BlockId with ShuffleId {
   override def name: String = RemoteBlockPushResolver.MERGED_SHUFFLE_FILE_NAME_PREFIX + "_" +
     appId + "_" + shuffleId + "_" + shuffleMergeId + "_" + reduceId + ".meta"
 }